from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

app = FastAPI(
    title="SQLCoder API",
    version="1.0.0"
)

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

MODEL_NAME = "defog/sqlcoder-7b-2"

print("Loading model...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
    trust_remote_code=True
)

print("Model loaded")


class SQLGenerationRequest(BaseModel):
    question: str
    table_name: str
    columns: list[str]
    dialect: str = "postgresql"


class SQLGenerationResponse(BaseModel):
    generated_sql: str


def build_prompt(
    question: str,
    table_name: str,
    columns: list[str],
    dialect: str
) -> str:

    schema = f"{table_name}({', '.join(columns)})"

    return f"""
### Task
Generate a {dialect} SQL query to answer the question.

### Database Schema
{schema}

### Instructions
- Use all requirements mentioned in the question.
- Use only columns from the schema.
- Return only SQL.
- Do not explain.
- Do not use markdown.

### Question
{question}

### SQL
"""


def generate_sql(
    question: str,
    table_name: str,
    columns: list[str],
    dialect: str
) -> str:

    prompt = build_prompt(
        question,
        table_name,
        columns,
        dialect
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt"
    ).to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.1,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id
    )

    generated_text = tokenizer.decode(
        outputs[0],
        skip_special_tokens=True
    )

    sql = generated_text.split("### SQL")[-1].strip()

    return sql


@app.get("/")
def root():
    return {"status": "running"}


@app.get("/health")
def health():
    return {
        "status": "healthy",
        "model": MODEL_NAME
    }


@app.post(
    "/generate-sql",
    response_model=SQLGenerationResponse
)
def generate_sql_endpoint(
    request: SQLGenerationRequest
):
    try:

        sql = generate_sql(
            question=request.question,
            table_name=request.table_name,
            columns=request.columns,
            dialect=request.dialect
        )

        return SQLGenerationResponse(
            generated_sql=sql
        )

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=str(e)
        )