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.0):
    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 if temperature > 0 else 1.0,
            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)}"

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: 2rem;">
            <h2 style="font-size: 2rem; font-weight: 700; margin-bottom: 0.5rem; color: white;">Interactive SQL Generation</h2>
            <p style="font-size: 1rem; color: #D0D0D0; margin: 0;">Transform natural language into 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"
            )
            
            gr.HTML("""
                <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0;">
                    <h4 style="color: #FF6B4A; font-size: 1rem; margin: 0 0 1.5rem 0; font-weight: 700;">Model Guide</h4>
                    
                    <div style="margin-bottom: 1rem;">
                        <div style="font-weight: 600; font-size: 0.9rem; margin-bottom: 0.4rem; color: #FFFFFF;">BM1 (33M parameters)</div>
                        <div style="font-size: 0.85rem; margin-left: 1rem; color: #FFFFFF;">TinyStories 33M fine-tuned</div>
                    </div>
                    
                    <div style="margin-bottom: 1rem;">
                        <div style="font-weight: 600; font-size: 0.9rem; margin-bottom: 0.4rem; color: #FFFFFF;">BM2 (0.5B parameters)</div>
                        <div style="font-size: 0.85rem; margin-left: 1rem; color: #FFFFFF;">Qwen 2.5 0.5B fine-tuned</div>
                    </div>
                    
                    <div style="margin-bottom: 1.5rem;">
                        <div style="font-weight: 600; font-size: 0.9rem; margin-bottom: 0.4rem; color: #FFFFFF;">BM3 (1B parameters)</div>
                        <div style="font-size: 0.85rem; margin-left: 1rem; color: #FFFFFF;">Llama 3.2 1B fine-tuned</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="font-size: 0.9rem; line-height: 2;">
                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS1:</strong> <span style="color: #FFFFFF;">Basic SELECT-FROM</span></div>
                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS2:</strong> <span style="color: #FFFFFF;">Adds ORDER BY</span></div>
                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS3:</strong> <span style="color: #FFFFFF;">Aggregations</span></div>
                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS4:</strong> <span style="color: #FFFFFF;">WHERE filters</span></div>
                            <div style="color: #FFFFFF;"><strong style="color: #FFFFFF;">CS5:</strong> <span style="color: #FFFFFF;">Multi-table JOINs</span></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,
                value=""
            )
            
            schema = gr.Code(
                label="Database Schema (SQL)",
                language="sql",
                value="""CREATE TABLE employees (
    name VARCHAR(100),
    salary INT,
    department VARCHAR(100)
);""",
                lines=6
            )
            
            with gr.Row():
                max_length = gr.Slider(64, 512, value=256, step=32, label="Max Length")
                temperature = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature")
            
            generate_btn = gr.Button("Generate SQL", variant="primary", size="lg")
            
            output = gr.Code(
                label="Generated SQL Query",
                language="sql",
                lines=8
            )

    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}