import streamlit as st
from pptx import Presentation
from pptx.util import Inches, Pt
from pptx.dml.color import RGBColor
from pptx.enum.text import PP_ALIGN, PP_PARAGRAPH_ALIGNMENT
import io
import tempfile
import os
import requests
import re

# Initialize Streamlit
st.set_page_config(page_title="PPTX Smart Enhancer", layout="wide")

# Configuration
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
MODEL_NAME = "gemma2-9b-it"

# Style Options
FONT_CHOICES = {
    "Arial": "Arial",
    "Calibri": "Calibri",
    "Times New Roman": "Times New Roman",
    "Helvetica": "Helvetica",
    "Verdana": "Verdana",
    "Georgia": "Georgia",
    "Courier New": "Courier New",
    "Palatino": "Palatino",
    "Garamond": "Garamond",
    "Trebuchet MS": "Trebuchet MS"
}

COLOR_PALETTES = {
    "Professional": {"primary": "#2B579A", "secondary": "#5B9BD5", "text": "#000000", "background": "#FFFFFF"},
    "Modern": {"primary": "#404040", "secondary": "#A5A5A5", "text": "#FFFFFF", "background": "#121212"},
    "Vibrant": {"primary": "#E53935", "secondary": "#FFCDD2", "text": "#212121", "background": "#F5F5F5"},
    "Corporate": {"primary": "#1976D2", "secondary": "#BBDEFB", "text": "#212121", "background": "#FAFAFA"},
    "Creative": {"primary": "#7B1FA2", "secondary": "#CE93D8", "text": "#FFFFFF", "background": "#212121"},
    "Nature": {"primary": "#388E3C", "secondary": "#A5D6A7", "text": "#212121", "background": "#F1F8E9"},
    "Elegant": {"primary": "#5D4037", "secondary": "#BCAAA4", "text": "#FFFFFF", "background": "#3E2723"},
    "Minimal": {"primary": "#000000", "secondary": "#E0E0E0", "text": "#212121", "background": "#FFFFFF"}
}

def get_custom_colors():
    """Get custom colors from user input with background option"""
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        primary = st.color_picker("Primary Color", "#2B579A")
    with col2:
        secondary = st.color_picker("Secondary Color", "#5B9BD5")
    with col3:
        text = st.color_picker("Text Color", "#000000")
    with col4:
        background = st.color_picker("Background Color", "#FFFFFF")
    return {
        "primary": primary,
        "secondary": secondary,
        "text": text,
        "background": background
    }

def get_ai_response(prompt: str) -> str:
    """Get formatted content from Groq API."""
    headers = {"Authorization": f"Bearer {GROQ_API_KEY}"}
    payload = {
        "messages": [{
            "role": "user",
            "content": f"{prompt}\n\nReturn well-structured content with clear headings, concise bullet points (max 5 per slide), and no markdown. Use bold/italic via formatting, not symbols."
        }],
        "model": MODEL_NAME,
        "temperature": 0.7
    }
    response = requests.post(
        "https://api.groq.com/openai/v1/chat/completions",
        headers=headers,
        json=payload,
        timeout=30
    )
    return response.json()["choices"][0]["message"]["content"]

def extract_slide_text(slide) -> str:
    """Extract text from slide."""
    return "\n".join(
        paragraph.text
        for shape in slide.shapes
        if hasattr(shape, "text_frame")
        for paragraph in shape.text_frame.paragraphs
    )

def split_formatted_runs(text):
    """Split text into formatted runs removing markdown."""
    formatted_segments = []
    bold_parts = re.split(r'(\*\*)', text)
    bold = False
    current_text = []
    
    for part in bold_parts:
        if part == '**':
            if current_text:
                formatted_segments.append({'text': ''.join(current_text), 'bold': bold, 'italic': False})
                current_text = []
            bold = not bold
        else:
            italic_parts = re.split(r'(\*)', part)
            italic = False
            for italic_part in italic_parts:
                if italic_part == '*':
                    if current_text:
                        formatted_segments.append({'text': ''.join(current_text), 'bold': bold, 'italic': italic})
                        current_text = []
                    italic = not italic
                else:
                    current_text.append(italic_part)
            if current_text:
                formatted_segments.append({'text': ''.join(current_text), 'bold': bold, 'italic': italic})
                current_text = []
    return formatted_segments

def apply_formatting(text_frame, content, design_settings):
    """Apply proper formatting to text frame with markdown removal."""
    text_frame.clear()
    paragraphs = [p for p in content.split('\n') if p.strip()]
    
    for para in paragraphs:
        clean_para = re.sub(r'^#+\s*', '', para).strip()
        
        if re.match(r'^(What is|Understanding|What\'s|Drive Theory|Incentive Theory)', clean_para, re.IGNORECASE):
            p = text_frame.add_paragraph()
            p.text = clean_para
            p.font.size = Pt(design_settings["heading_size"])
            p.font.bold = True
            p.space_after = Pt(12)
            continue
            
        if re.match(r'^[\•\-\*] ', clean_para):
            p = text_frame.add_paragraph()
            p.level = 0
            content = clean_para[2:].strip()
            segments = split_formatted_runs(content)
            for seg in segments:
                run = p.add_run()
                run.text = seg['text']
                run.font.bold = seg['bold']
                run.font.italic = seg['italic']
            p.font.size = Pt(design_settings["body_size"])
            continue
            
        p = text_frame.add_paragraph()
        segments = split_formatted_runs(clean_para)
        for seg in segments:
            run = p.add_run()
            run.text = seg['text']
            run.font.bold = seg['bold']
            run.font.italic = seg['italic']
        p.font.size = Pt(design_settings["body_size"])
        p.space_after = Pt(6)

def enhance_slide(slide, user_prompt, design_settings):
    """Enhance slide content with proper formatting."""
    original_text = extract_slide_text(slide)
    if not original_text.strip():
        return
    
    prompt = f"Improve this slide content:\n{original_text}\n\nInstructions: {user_prompt}"
    enhanced_content = get_ai_response(prompt)
    
    title_shape = getattr(slide.shapes, 'title', None)
    shapes_to_remove = [s for s in slide.shapes if s != title_shape]
    
    for shape in shapes_to_remove:
        sp = shape._element
        sp.getparent().remove(sp)
    
    left = Inches(1) if title_shape else Inches(0.5)
    top = Inches(1.5) if title_shape else Inches(0.5)
    textbox = slide.shapes.add_textbox(left, top, Inches(8), Inches(5))
    text_frame = textbox.text_frame
    text_frame.word_wrap = True
    
    apply_formatting(text_frame, enhanced_content, design_settings)
    apply_design(slide, design_settings)

def apply_design(slide, design_settings):
    """Apply visual design to slide."""
    colors = design_settings["colors"]
    
    for shape in slide.shapes:
        if hasattr(shape, "text_frame"):
            for paragraph in shape.text_frame.paragraphs:
                for run in paragraph.runs:
                    run.font.name = design_settings["font"]
                    run.font.color.rgb = RGBColor.from_string(colors["text"][1:])
                    
                    if shape == getattr(slide.shapes, 'title', None):
                        run.font.color.rgb = RGBColor.from_string(colors["primary"][1:])
                        run.font.size = Pt(design_settings["title_size"])
                        run.font.bold = True
    
    if design_settings["set_background"]:
        slide.background.fill.solid()
        slide.background.fill.fore_color.rgb = RGBColor.from_string(colors["background"][1:])

def process_presentation(uploaded_file, user_prompt, design_settings):
    """Process and enhance the entire presentation."""
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pptx") as tmp_file:
        tmp_file.write(uploaded_file.read())
        tmp_path = tmp_file.name
    
    prs = Presentation(tmp_path)
    
    with st.spinner("Enhancing slides..."):
        for slide in prs.slides:
            enhance_slide(slide, user_prompt, design_settings)
    
    os.unlink(tmp_path)
    return prs

def main():
    st.title("Professional PPTX Enhancer")
    
    uploaded_file = st.file_uploader("Upload PowerPoint", type=["pptx"])
    user_prompt = st.text_area(
        "Enhancement Instructions",
        placeholder="Example: 'Make headings clear and add bullet points for key items'",
        height=100
    )
    
    with st.expander("🎨 Design Options", expanded=True):
        tab1, tab2 = st.tabs(["Preset Themes", "Custom Colors"])
        
        with tab1:
            col1, col2 = st.columns(2)
            with col1:
                font = st.selectbox("Font", list(FONT_CHOICES.keys()), index=1)
                color_palette = st.selectbox("Color Scheme", list(COLOR_PALETTES.keys()))
            with col2:
                title_size = st.slider("Title Size", 12, 44, 32)
                heading_size = st.slider("Heading Size", 12, 36, 24)
                body_size = st.slider("Body Text Size", 8, 28, 18)
                set_background = st.checkbox("Apply Background", True)
        
        with tab2:
            st.write("Create your own color scheme:")
            custom_colors = get_custom_colors()
            use_custom = st.checkbox("Use Custom Colors")
    
    design_settings = {
        "font": FONT_CHOICES[font],
        "colors": custom_colors if use_custom else COLOR_PALETTES[color_palette],
        "color_palette": "Custom" if use_custom else color_palette,
        "title_size": title_size,
        "heading_size": heading_size,
        "body_size": body_size,
        "set_background": set_background
    }
    
    if st.button("Enhance Presentation") and uploaded_file and user_prompt:
        try:
            if not GROQ_API_KEY:
                st.error("GROQ API key not found! Please set it in Space settings.")
                st.stop()
                
            enhanced_prs = process_presentation(uploaded_file, user_prompt, design_settings)
            
            buffer = io.BytesIO()
            enhanced_prs.save(buffer)
            buffer.seek(0)
            
            st.success("Done!")
            st.download_button(
                "Download Enhanced PPTX",
                data=buffer,
                file_name="enhanced.pptx",
                mime="application/vnd.openxmlformats-officedocument.presentationml.presentation"
            )
        except Exception as e:
            st.error(f"Error: {str(e)}")

if __name__ == "__main__":
    main()