import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODELS = {
    "BM1_CS1_Syn (33M)": "withmartian/sql_interp_bm1_cs1_experiment_1.10",
    "BM1_CS2_Syn (33M)": "withmartian/sql_interp_bm1_cs2_experiment_2.10",
    "BM1_CS3_Syn (33M)": "withmartian/sql_interp_bm1_cs3_experiment_3.10",
    "BM1_CS4_Syn (33M)": "withmartian/sql_interp_bm1_cs4_dataset_synonyms_experiment_1.1",
    "BM1_CS5_Syn (33M)": "withmartian/sql_interp_bm1_cs5_dataset_synonyms_experiment_1.2",
    "BM2_CS1_Syn (0.5B)": "withmartian/sql_interp_bm2_cs1_experiment_4.3",
    "BM2_CS2_Syn (0.5B)": "withmartian/sql_interp_bm2_cs2_experiment_5.3",
    "BM2_CS3_Syn (0.5B)": "withmartian/sql_interp_bm2_cs3_experiment_6.3",
    "BM3_CS1_Syn (1B)": "withmartian/sql_interp_bm3_cs1_experiment_7.3",
    "BM3_CS2_Syn (1B)": "withmartian/sql_interp_bm3_cs2_experiment_8.3",
    "BM3_CS3_Syn (1B)": "withmartian/sql_interp_bm3_cs3_experiment_9.3",
}

model_cache = {}

def load_model(model_name):
    if model_name not in model_cache:
        model_id = MODELS[model_name]
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        model_cache[model_name] = (tokenizer, model)
    return model_cache[model_name]

def generate_sql(model_name, instruction, schema, max_length=256, temperature=0.7):
    if not model_name or not instruction or not schema:
        return "Please fill in all fields and select a model"
    
    try:
        tokenizer, model = load_model(model_name)
        
        prompt = f"""### Instruction: {instruction}
### Context: {schema}
### Response:"""
        
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            temperature=temperature,
            do_sample=temperature > 0,
            pad_token_id=tokenizer.eos_token_id
        )
        
        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
        if "### Response:" in generated:
            sql = generated.split("### Response:")[-1].strip()
        else:
            sql = generated.strip()
        return sql
        
    except Exception as e:
        return f"Error: {str(e)}"

# Only ONE example
examples = [
    [
        "BM2_CS2_Syn (0.5B)",
        "List worker earnings from highest to lowest",
        "CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))"
    ],
]

def model_demo(shared_instruction, shared_schema):
    gr.HTML("""
        <div style="text-align: center; padding: 2rem; background: linear-gradient(135deg, #3A3A3A 0%, #4A4A4A 100%); border-radius: 16px; margin-bottom: 1.5rem;">
            <h2 style="font-size: 2rem; font-weight: 700; margin-bottom: 0.5rem; color: white;">Interactive SQL Generation</h2>
            <p style="font-size: 1rem; opacity: 0.9; line-height: 1.6; color: #D0D0D0;">
                Transform natural language into SQL using <span style="color: #FF6B4A; font-weight: 600;">mechanistically interpretable</span> models
            </p>
        </div>
    """)
    
    gr.HTML("""
        <div style="background: linear-gradient(135deg, #2A2A2A 0%, #3A3A3A 100%); border-radius: 12px; padding: 1.5rem; margin: 1rem 0; border-left: 4px solid #FF6B4A;">
            <p style="color: #D0D0D0; margin: 0; line-height: 1.6;">
                <strong style="color: #FF6B4A;">How it works:</strong> Select a model, describe your query in plain English, and watch the model generate SQL.
            </p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Configuration")
            model_dropdown = gr.Dropdown(
                choices=list(MODELS.keys()),
                value="BM2_CS2_Syn (0.5B)",
                label="Model Selection",
                info="Larger models = better accuracy"
            )
            
            # Better formatted guide
            gr.HTML("""
                <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0; border: 1px solid #3A3A3A;">
                    <h4 style="color: #FF6B4A; font-size: 1rem; margin: 0 0 1.25rem 0; font-weight: 700; border-bottom: 2px solid #3A3A3A; padding-bottom: 0.75rem;">Model Guide</h4>
                    
                    <div style="margin-bottom: 1rem;">
                        <div style="color: #4CAF50; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.35rem;">BM1 (33M parameters)</div>
                        <div style="color: #999; font-size: 0.85rem; margin-left: 1rem;">Lightning fast, ideal for simple queries</div>
                    </div>
                    
                    <div style="margin-bottom: 1rem;">
                        <div style="color: #FF9800; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.35rem;">BM2 (0.5B parameters)</div>
                        <div style="color: #999; font-size: 0.85rem; margin-left: 1rem;">Balanced performance and speed</div>
                    </div>
                    
                    <div style="margin-bottom: 1.25rem;">
                        <div style="color: #f44336; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.35rem;">BM3 (1B parameters)</div>
                        <div style="color: #999; font-size: 0.85rem; margin-left: 1rem;">Most accurate for complex queries</div>
                    </div>
                    
                    <div style="padding-top: 1rem; border-top: 1px solid #3A3A3A;">
                        <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.75rem;">Dataset Complexity</div>
                        <div style="color: #999; font-size: 0.85rem; line-height: 1.8;">
                            <div><strong style="color: #D0D0D0;">CS1:</strong> Basic SELECT-FROM</div>
                            <div><strong style="color: #D0D0D0;">CS2:</strong> Adds ORDER BY</div>
                            <div><strong style="color: #D0D0D0;">CS3:</strong> Aggregations</div>
                            <div><strong style="color: #D0D0D0;">CS4:</strong> WHERE filters</div>
                            <div><strong style="color: #D0D0D0;">CS5:</strong> Multi-table JOINs</div>
                        </div>
                    </div>
                </div>
            """)

        with gr.Column(scale=2):
            gr.Markdown("### Your Query")
            instruction = gr.Textbox(
                label="Natural Language Query",
                placeholder="e.g., Find all employees earning more than 50000 sorted by name",
                lines=2
            )
            
            # Database schema as Code (SQL formatted)
            schema = gr.Code(
                label="Database Schema",
                language="sql",
                value="CREATE TABLE employees (name VARCHAR(100), salary INT, department VARCHAR(100))",
                lines=3
            )
            
            with gr.Row():
                max_length = gr.Slider(64, 512, value=256, step=32, label="Max Length", info="Token limit")
                temperature = gr.Slider(0.0, 1.0, value=0.1, step=0.1, label="Temperature", info="Creativity level")
            
            generate_btn = gr.Button("Generate SQL", variant="primary", size="lg")
            
            # Output as Code (SQL formatted)
            output = gr.Code(label="Generated SQL Query", language="sql", lines=6)

    gr.Markdown("### Example Query")
    
    # Custom examples with 2-line schema display
    gr.HTML("""
        <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1rem 0; border: 1px solid #3A3A3A;">
            <div style="margin-bottom: 0.75rem;">
                <strong style="color: #FF6B4A;">Model:</strong> <span style="color: #D0D0D0;">BM2_CS2_Syn (0.5B)</span>
            </div>
            <div style="margin-bottom: 0.75rem;">
                <strong style="color: #FF6B4A;">Query:</strong> <span style="color: #D0D0D0;">List worker earnings from highest to lowest</span>
            </div>
            <div>
                <strong style="color: #FF6B4A;">Schema:</strong>
                <pre style="background: #1A1A1A; padding: 0.75rem; border-radius: 6px; margin-top: 0.5rem; color: #D0D0D0; font-size: 0.85rem; line-height: 1.5;">CREATE TABLE employees (
    name VARCHAR(100), salary INT, department VARCHAR(100)
)</pre>
            </div>
        </div>
    """)

    shared_instruction.change(
        fn=lambda x: x,
        inputs=shared_instruction,
        outputs=instruction
    )
    
    shared_schema.change(
        fn=lambda x: x,
        inputs=shared_schema,
        outputs=schema
    )

    generate_btn.click(
        fn=generate_sql,
        inputs=[model_dropdown, instruction, schema, max_length, temperature],
        outputs=output
    )
    
    return {'instruction': instruction, 'schema': schema, 'output': output}