import gradio as gr
import json
import os
from PIL import Image
import tempfile
from utils.extractor import DesignTokenExtractor
from utils.token_generator import TokenCodeGenerator


def create_token_preview(tokens):
    """Create HTML preview of extracted tokens"""
    html = """
    <div style="font-family: system-ui, sans-serif; padding: 20px; background: #f9fafb; border-radius: 8px;">
        <h3 style="margin-top: 0; color: #1f2937;">Extracted Design Tokens</h3>
    """
    
    # Color palette preview
    if 'colors' in tokens and tokens['colors']:
        html += """
        <div style="margin-bottom: 24px;">
            <h4 style="color: #6b7280; font-size: 14px; text-transform: uppercase; letter-spacing: 0.05em;">Colors</h4>
            <div style="display: flex; gap: 12px; flex-wrap: wrap;">
        """
        for name, color in tokens['colors'].items():
            html += f"""
            <div style="text-align: center;">
                <div style="width: 80px; height: 80px; background: {color['hex']}; 
                           border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);"></div>
                <div style="margin-top: 8px;">
                    <div style="font-size: 12px; font-weight: 600; color: #374151;">{name}</div>
                    <div style="font-size: 11px; color: #9ca3af;">{color['hex']}</div>
                    <div style="font-size: 10px; color: #9ca3af;">{int(color.get('proportion', 0) * 100)}%</div>
                </div>
            </div>
            """
        html += "</div></div>"
    
    # Spacing preview
    if 'spacing' in tokens and tokens['spacing']:
        html += """
        <div style="margin-bottom: 24px;">
            <h4 style="color: #6b7280; font-size: 14px; text-transform: uppercase; letter-spacing: 0.05em;">Spacing</h4>
            <div style="display: flex; gap: 16px; align-items: flex-end;">
        """
        for name, value in tokens['spacing'].items():
            try:
                height = value.replace('px', '')
                html += f"""
                <div style="text-align: center;">
                    <div style="width: 60px; height: {height}px; background: #3b82f6; 
                               border-radius: 4px; opacity: 0.8;"></div>
                    <div style="margin-top: 8px;">
                        <div style="font-size: 12px; font-weight: 600; color: #374151;">{name}</div>
                        <div style="font-size: 11px; color: #9ca3af;">{value}</div>
                    </div>
                </div>
                """
            except:
                pass
        html += "</div></div>"
    
    # Typography preview
    if 'typography' in tokens and tokens['typography']:
        html += """
        <div style="margin-bottom: 24px;">
            <h4 style="color: #6b7280; font-size: 14px; text-transform: uppercase; letter-spacing: 0.05em;">Typography</h4>
        """
        for name, props in tokens['typography'].items():
            size = props.get('size', '16px')
            weight = props.get('weight', '400')
            family = props.get('family', 'sans-serif')
            html += f"""
            <div style="margin-bottom: 12px; padding: 12px; background: white; border-radius: 6px;">
                <div style="font-size: {size}; font-weight: {weight}; font-family: {family}; color: #1f2937;">
                    Sample {name.title()} Text
                </div>
                <div style="font-size: 11px; color: #9ca3af; margin-top: 4px;">
                    {family} • {size} • Weight {weight}
                </div>
            </div>
            """
        html += "</div>"
    
    html += "</div>"
    return html


def process_screenshot(image, output_format, progress=None):
    """Process uploaded screenshot and extract design tokens"""
    if image is None:
        return None, "Please upload a screenshot", None
    
    extractor = DesignTokenExtractor()
    generator = TokenCodeGenerator()
    
    try:
        if progress:
            progress(0.1, desc="Initializing extraction...")
        
        # Resize image if needed
        image = extractor.resize_for_processing(image)
        
        # Save temporary file for colorgram
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
            temp_path = tmp.name
            image.save(temp_path)
        
        if progress:
            progress(0.3, desc="Extracting colors...")
        colors = extractor.extract_colors(temp_path)
        
        if progress:
            progress(0.5, desc="Detecting spacing...")
        spacing = extractor.detect_spacing(image)
        
        if progress:
            progress(0.6, desc="Analyzing typography...")
        typography = extractor.detect_typography(image)
        
        if progress:
            progress(0.7, desc="Analyzing components...")
        components = extractor.analyze_components(image)
        
        # Combine all tokens
        tokens = {
            "colors": colors,
            "spacing": spacing,
            "typography": typography,
            "components": components
        }
        
        if progress:
            progress(0.8, desc="Generating code...")
        
        # Generate output based on selected format
        if output_format == "CSS Variables":
            code_output = generator.generate_css_variables(tokens)
            file_ext = "css"
        elif output_format == "Tailwind Config":
            code_output = generator.generate_tailwind_config(tokens)
            file_ext = "js"
        elif output_format == "JSON Tokens":
            code_output = generator.generate_json_tokens(tokens)
            file_ext = "json"
        elif output_format == "Style Dictionary":
            code_output = generator.generate_style_dictionary(tokens)
            file_ext = "json"
        elif output_format == "SCSS Variables":
            code_output = generator.generate_scss_variables(tokens)
            file_ext = "scss"
        else:
            code_output = json.dumps(tokens, indent=2)
            file_ext = "json"
        
        # Save output file
        output_filename = f"design_tokens.{file_ext}"
        with open(output_filename, "w") as f:
            f.write(code_output)
        
        # Clean up temp file
        try:
            os.unlink(temp_path)
        except:
            pass
        
        if progress:
            progress(1.0, desc="Complete!")
        
        # Create preview visualization
        preview_html = create_token_preview(tokens)
        
        return preview_html, code_output, output_filename
        
    except Exception as e:
        return None, f"Error processing screenshot: {str(e)}", None


def create_gradio_app():
    """Create the main Gradio application"""
    
    with gr.Blocks(
        title="Design Token Extractor",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            font-family: 'Inter', system-ui, sans-serif;
        }
        .gr-button-primary {
            background-color: #3b82f6 !important;
        }
        """
    ) as app:
        gr.Markdown(
            """
            # 🎨 Design Token Extractor
            
            Transform UI screenshots into structured design token libraries using AI-powered analysis.
            Upload a screenshot to automatically extract colors, spacing, typography, and component tokens.
            
            ---
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="Upload UI Screenshot",
                    type="pil",
                    sources=["upload", "clipboard"],
                    height=400
                )
                
                output_format = gr.Radio(
                    choices=[
                        "CSS Variables",
                        "Tailwind Config",
                        "JSON Tokens",
                        "Style Dictionary",
                        "SCSS Variables"
                    ],
                    value="CSS Variables",
                    label="Output Format",
                    info="Choose the format for your design tokens"
                )
                
                extract_btn = gr.Button(
                    "🚀 Extract Design Tokens",
                    variant="primary",
                    size="lg"
                )
                
                gr.Markdown(
                    """
                    ### Tips for best results:
                    - Use high-quality screenshots (min 800px width)
                    - Include various UI elements for comprehensive extraction
                    - Screenshots with clear color hierarchy work best
                    - Ensure good contrast between elements
                    """
                )
            
            with gr.Column(scale=1):
                preview = gr.HTML(
                    label="Token Preview",
                    value="<div style='padding: 20px; text-align: center; color: #9ca3af;'>Upload a screenshot to see extracted tokens</div>"
                )
                
                code_output = gr.Code(
                    label="Generated Code",
                    language="css",
                    lines=20,
                    value="// Your design tokens will appear here"
                )
                
                download_file = gr.File(
                    label="Download Tokens",
                    visible=True
                )
        
        # Examples section commented out until example images are available
        # gr.Markdown("### Example Screenshots")
        # gr.Examples(
        #     examples=[
        #         ["examples/dashboard.png", "CSS Variables"],
        #         ["examples/landing_page.png", "Tailwind Config"],
        #         ["examples/mobile_app.png", "JSON Tokens"]
        #     ],
        #     inputs=[input_image, output_format],
        #     cache_examples=False
        # )
        
        # Connect the extraction function
        extract_btn.click(
            fn=process_screenshot,
            inputs=[input_image, output_format],
            outputs=[preview, code_output, download_file],
            show_progress=True
        )
        
        # Add footer
        gr.Markdown(
            """
            ---
            
            ### Features:
            - **Color Extraction**: Identifies dominant colors and creates semantic color roles
            - **Spacing Detection**: Analyzes layout patterns to extract consistent spacing values
            - **Typography Analysis**: Detects font styles and creates text hierarchy tokens
            - **Multiple Output Formats**: Export to CSS, Tailwind, JSON, Style Dictionary, or SCSS
            
            Built with ❤️ using Gradio and computer vision models
            """
        )
    
    return app


if __name__ == "__main__":
    app = create_gradio_app()
    app.launch(
        server_name="0.0.0.0", 
        server_port=7860,
        share=True,
        show_error=True
    )