#!/usr/bin/env python3

import os
import sys
import tempfile
import zipfile
import json
from pathlib import Path
from typing import Optional, Tuple
import uuid
from datetime import datetime
import traceback

import gradio as gr

project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

try:
    from src.state.poster_state import create_state
    from src.workflow.pipeline import create_workflow_graph
except ImportError as e:
    print(f"Error importing modules: {e}")
    sys.exit(1)

def set_temp_api_keys(anthropic_key, openai_key, anthropic_base_url=None, openai_base_url=None):
    """Temporarily set API keys and base URLs in environment, returns cleanup function"""
    original_values = {}
    
    # Save original values and set new ones
    if anthropic_key and anthropic_key.strip():
        original_values["ANTHROPIC_API_KEY"] = os.environ.get("ANTHROPIC_API_KEY")
        os.environ["ANTHROPIC_API_KEY"] = anthropic_key.strip()
    if openai_key and openai_key.strip():
        original_values["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY")
        os.environ["OPENAI_API_KEY"] = openai_key.strip()
    if anthropic_base_url and anthropic_base_url.strip():
        original_values["ANTHROPIC_BASE_URL"] = os.environ.get("ANTHROPIC_BASE_URL")
        os.environ["ANTHROPIC_BASE_URL"] = anthropic_base_url.strip()
    if openai_base_url and openai_base_url.strip():
        original_values["OPENAI_BASE_URL"] = os.environ.get("OPENAI_BASE_URL")
        os.environ["OPENAI_BASE_URL"] = openai_base_url.strip()
    
    def cleanup():
        """Restore original environment values"""
        for key, original_value in original_values.items():
            if original_value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = original_value
    
    return cleanup

AVAILABLE_MODELS = [
    "claude-sonnet-4-20250514",
    "gpt-4o-2024-08-06", 
    "gpt-4.1-2025-04-14",
    "gpt-4.1-mini-2025-04-14"
]

def create_job_directory() -> Path:
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    job_id = str(uuid.uuid4())[:8]
    dir_name = f"job_{timestamp}_{job_id}"
    job_dir = Path(tempfile.mkdtemp(prefix=f"{dir_name}_"))
    return job_dir

def validate_inputs(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, anthropic_base_url, openai_base_url):
    if not pdf_file:
        return "Please upload PDF paper"
    if not logo_file:
        return "Please upload conference logo"
    if not aff_logo_file:
        return "Please upload affiliation logo"
    if text_model not in AVAILABLE_MODELS:
        return f"Invalid text model: {text_model}"
    if vision_model not in AVAILABLE_MODELS:
        return f"Invalid vision model: {vision_model}"
    
    # Check API keys
    has_anthropic = bool(anthropic_key and anthropic_key.strip())
    has_openai = bool(openai_key and openai_key.strip())
    
    if not has_anthropic and not has_openai:
        return "Please provide at least one API key (Anthropic or OpenAI)"
    
    # Check if selected models have corresponding API keys
    if text_model.startswith("claude") and not has_anthropic:
        return "Anthropic API key required for Claude models"
    if text_model.startswith("gpt") and not has_openai:
        return "OpenAI API key required for GPT models"
    if vision_model.startswith("claude") and not has_anthropic:
        return "Anthropic API key required for Claude models"
    if vision_model.startswith("gpt") and not has_openai:
        return "OpenAI API key required for GPT models"
    
    ratio = poster_width / poster_height
    if ratio < 1.4 or ratio > 2.0:
        return f"Poster ratio {ratio:.2f} out of range (1.4-2.0)"
    
    # Check file type - Gradio returns file object with name attribute
    if hasattr(pdf_file, 'name') and not pdf_file.name.lower().endswith('.pdf'):
        return "Paper must be PDF format"
    
    return None

def generate_poster(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, anthropic_base_url, openai_base_url, progress=gr.Progress()):
    try:
        # Set API keys temporarily
        cleanup_api_keys = set_temp_api_keys(anthropic_key, openai_key, anthropic_base_url, openai_base_url)
        
        error_msg = validate_inputs(pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, anthropic_base_url, openai_base_url)
        if error_msg:
            cleanup_api_keys()
            return None, f"❌ {error_msg}"
        
        progress(0.1, desc="Initializing...")
        
        job_dir = create_job_directory()
        
        pdf_path = job_dir / "paper.pdf"
        logo_path = job_dir / "logo.png" 
        aff_logo_path = job_dir / "aff_logo.png"
        
        # Handle file writing - check if it's file object or bytes
        if hasattr(pdf_file, 'read'):
            pdf_content = pdf_file.read()
        else:
            pdf_content = pdf_file
        
        if hasattr(logo_file, 'read'):
            logo_content = logo_file.read()
        else:
            logo_content = logo_file
            
        if hasattr(aff_logo_file, 'read'):
            aff_logo_content = aff_logo_file.read()
        else:
            aff_logo_content = aff_logo_file
        
        with open(pdf_path, "wb") as f:
            f.write(pdf_content)
        with open(logo_path, "wb") as f:
            f.write(logo_content)
        with open(aff_logo_path, "wb") as f:
            f.write(aff_logo_content)
            
        progress(0.2, desc="Setting up workflow...")
        
        state = create_state(
            pdf_path=str(pdf_path),
            text_model=text_model,
            vision_model=vision_model,
            width=int(poster_width),
            height=int(poster_height),
            url="",
            logo_path=str(logo_path),
            aff_logo_path=str(aff_logo_path)
        )
        
        progress(0.3, desc="Compiling workflow...")
        
        graph = create_workflow_graph()
        workflow = graph.compile()
        
        progress(0.5, desc="Processing paper...")
        
        final_state = workflow.invoke(state)
        
        progress(0.8, desc="Generating outputs...")
        
        if final_state.get("errors"):
            error_details = "; ".join(final_state["errors"])
            cleanup_api_keys()
            return None, f"❌ Generation errors: {error_details}"
        
        output_dir = Path(final_state["output_dir"])
        poster_name = final_state.get("poster_name", "poster")
        
        progress(0.9, desc="Packaging results...")
        
        zip_path = job_dir / f"{poster_name}_output.zip"
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            if output_dir.exists():
                for file_path in output_dir.rglob("*"):
                    if file_path.is_file():
                        arcname = file_path.relative_to(output_dir)
                        zipf.write(file_path, arcname)
        
        
        progress(1.0, desc="✅ Completed!")
        
        success_msg = f"""✅ Poster generation successful!

Poster: {poster_name}
Output: {output_dir.name}
Package: {zip_path.name}
Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"""
        
        cleanup_api_keys()
        return str(zip_path), success_msg
        
    except Exception as e:
        cleanup_api_keys()
        error_msg = f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
        return None, error_msg


with gr.Blocks(title="PosterGen", css="""
    .gradio-column {
        margin-left: 10px !important;
        margin-right: 10px !important;
    }
    .gradio-column:first-child {
        margin-left: 0 !important;
    }
    .gradio-column:last-child {
        margin-right: 0 !important;
    }
    """) as demo:
    
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 30px;">
        <h1 style="margin-bottom: 10px;">
            PosterGen
        </h1>
        <p style="font-size: 18px; color: #666;">🎨 Generate design-aware academic posters from PDF papers</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1, variant="panel"):
            gr.Markdown("### 📄 Upload Files")
            
            pdf_file = gr.File(label="PDF Paper", file_types=[".pdf"], type="binary")
            
            with gr.Row():
                logo_file = gr.File(label="Conference Logo", file_types=["image"], type="binary")
                aff_logo_file = gr.File(label="Affiliation Logo", file_types=["image"], type="binary")
        
        with gr.Column(scale=1, variant="panel"):
            with gr.Group():
                gr.Markdown("### 🔑 API Keys")
                gr.Markdown("⚠️ Keys are processed securely and not stored")
                
                with gr.Row():
                    anthropic_key = gr.Textbox(
                        label="Anthropic API Key", 
                        type="password", 
                        placeholder="sk-ant-...",
                        info="Required for Claude models"
                    )
                    openai_key = gr.Textbox(
                        label="OpenAI API Key", 
                        type="password", 
                        placeholder="sk-...",
                        info="Required for GPT models"
                    )
                    
                with gr.Row():
                    anthropic_base_url = gr.Textbox(
                        label="Anthropic Base URL (Optional)",
                        placeholder="https://api.anthropic.com",
                        info="Set the base url for compatible API services"
                    )
                    openai_base_url = gr.Textbox(
                        label="OpenAI Base URL (Optional)",
                        placeholder="https://api.openai.com/v1",
                        info="Set the base url for compatible API services"
                    )
            
            gr.Markdown("### 🤖 Model Settings")
            
            with gr.Row():
                text_model = gr.Dropdown(choices=AVAILABLE_MODELS, value=AVAILABLE_MODELS[0], label="Text Model")
                vision_model = gr.Dropdown(choices=AVAILABLE_MODELS, value=AVAILABLE_MODELS[0], label="Vision Model")
            
            gr.Markdown("### 📐 Dimensions")
            
            with gr.Row():
                poster_width = gr.Number(value=54, minimum=20, maximum=100, step=0.1, label="Width (inches)")
                poster_height = gr.Number(value=36, minimum=10, maximum=60, step=0.1, label="Height (inches)")
            
            
        with gr.Column(scale=1, variant="panel"):
            gr.Markdown("### 📊 Status")
            status_output = gr.Textbox(label="Generation Status", placeholder="Click 'Generate Poster' to start...", lines=6)
            
            gr.Markdown("### 📥 Download")
            download_file = gr.File(label="Download Package")
    
    # Generate button spanning full width
    with gr.Row():
        generate_btn = gr.Button("🚀 Generate Poster", variant="primary", size="lg")
    
    def generate_and_display(*args):
        download_file_result, status_result = generate_poster(*args)
        return download_file_result, status_result
    
    generate_btn.click(
        fn=generate_and_display,
        inputs=[pdf_file, logo_file, aff_logo_file, text_model, vision_model, poster_width, poster_height, anthropic_key, openai_key, anthropic_base_url, openai_base_url],
        outputs=[download_file, status_output]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)