import json
import traceback
import gradio as gr
import spaces

from app.api.schemas import TableSchema, ForeignKey
from app.engine.validator import select_valid, format_schema
from app.engine.inference import generate_sql

def parse_schema_input(schema_raw: str):
    """
    Parses either JSON-serialized schemas or plain string inputs.
    Returns: (formatted_schema_str, tables_list, fks_list)
    """
    if not schema_raw or not schema_raw.strip():
        return "", [], []

    try:
        data = json.loads(schema_raw)
        if isinstance(data, dict):
            tables_data = data.get("tables", [])
            fks_data = data.get("foreign_keys", [])
            
            tables = [TableSchema(**t) for t in tables_data]
            fks = [ForeignKey(**f) for f in fks_data]
            return format_schema(tables, fks), tables, fks
    except Exception:
        pass

    return schema_raw.strip(), [], []

@spaces.GPU(duration=60)
def predict_sql(question: str, schema_input: str = ""):
    try:
        schema_text, tables, foreign_keys = parse_schema_input(schema_input)
        prompt = f"[SCHEMA] {schema_text} [QUESTION] {question} [SQL] "
        
        candidates = generate_sql(prompt, n_candidates=5)
        
        if not candidates:
            return {"output": "SELECT 1;", "compiles": False, "candidates": []}

        if tables:
            chosen_sql, compiles = select_valid(tables, foreign_keys, candidates)
        else:
            chosen_sql, compiles = select_valid([], [], candidates)

        return {
            "output": chosen_sql,
            "compiles": compiles,
            "candidates": candidates
        }

    except Exception as e:
        return {
            "output": "",
            "compiles": False,
            "error": str(e),
            "traceback": traceback.format_exc()
        }

with gr.Blocks(title="Text-to-SQL Microservice") as demo:
    gr.Markdown("# Text-to-SQL Backend Microservice")
    
    with gr.Row():
        question_in = gr.Textbox(
            label="Question", 
            placeholder="What are the names of all employees?"
        )
        schema_in = gr.Textbox(
            label="Schema (Raw text or JSON with tables/foreign_keys)", 
            placeholder='{"tables": [{"name": "employees", "columns": ["id", "name"], "primary_keys": ["id"]}], "foreign_keys": []}'
        )
    
    json_out = gr.JSON(label="Result Payload")
    btn = gr.Button("Generate SQL")

    btn.click(
        fn=predict_sql,
        inputs=[question_in, schema_in],
        outputs=json_out,
        api_name="generate"
    )

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