""" Gradio UI for the Qwen2.5 Text-to-SQL LoRA model."""

from __future__ import annotations

import html
import os
import threading
from pathlib import Path
from typing import Any

# ZeroGPU bootstrap
IS_HF_SPACE = bool(os.getenv("SPACE_ID"))

try:
    import spaces
except ImportError:
    
    if IS_HF_SPACE:
        raise

    # Local-development fallback when the `spaces` package is not installed
    class _SpacesShim:
        @staticmethod
        def GPU(duration: int = 60, **_kwargs):
            def decorator(function):
                return function

            return decorator

    spaces = _SpacesShim()

import gradio as gr
import sqlglot
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

from config import MODEL_ID, OUTPUT_DIR, SYSTEM_PROMPT


# Runtime configuration
ADAPTER_ID = os.getenv("ADAPTER_ID", OUTPUT_DIR).strip()
HF_TOKEN = os.getenv("HF_TOKEN") or None

MODEL: Any | None = None
TOKENIZER: Any | None = None
MODEL_LOCK = threading.Lock()
MODEL_MODE = "not loaded"
MODEL_SOURCE = ""

DIALECT_MAP = {
    "Auto / Generic SQL": None,
    "SQLite": "sqlite",
    "PostgreSQL": "postgres",
    "MySQL": "mysql",
    "Microsoft SQL Server": "tsql",
}


# Visual design
CSS = r"""
:root {
  --surface: rgba(24, 19, 8, .80);
  --surface-2: rgba(38, 29, 8, .72);
  --surface-3: rgba(255, 255, 255, .035);
  --line: rgba(250, 204, 21, .17);
  --line-strong: rgba(250, 204, 21, .34);
  --muted: #b9ad8c;
  --text: #fffaf0;
  --accent: #facc15;
  --accent-2: #f59e0b;
  --accent-3: #fde68a;
  --success: #86efac;
  --danger: #fca5a5;
}

html, body {
  background: #090704 !important;
}

.gradio-container {
  max-width: 1500px !important;
  margin: 0 auto !important;
  color: var(--text) !important;
  background:
    radial-gradient(circle at 7% 7%, rgba(250, 204, 21, .20), transparent 30%),
    radial-gradient(circle at 91% 12%, rgba(245, 158, 11, .16), transparent 28%),
    radial-gradient(circle at 52% 92%, rgba(234, 179, 8, .08), transparent 33%),
    linear-gradient(145deg, #070603 0%, #100c04 46%, #171006 100%) !important;
  min-height: 100vh;
}

.main-shell {
  padding: 28px 24px 44px;
}

.hero {
  position: relative;
  overflow: hidden;
  border: 1px solid var(--line);
  background:
    linear-gradient(135deg, rgba(38, 29, 8, .95), rgba(15, 12, 6, .89));
  border-radius: 25px;
  padding: 31px 33px;
  box-shadow: 0 30px 85px rgba(0, 0, 0, .35);
  margin-bottom: 18px;
}

.hero::before {
  content: "";
  position: absolute;
  width: 390px;
  height: 390px;
  left: -185px;
  bottom: -275px;
  background: radial-gradient(circle, rgba(250, 204, 21, .20), transparent 66%);
}

.hero::after {
  content: "";
  position: absolute;
  width: 350px;
  height: 350px;
  right: -120px;
  top: -175px;
  background: radial-gradient(circle, rgba(245, 158, 11, .25), transparent 66%);
}

.eyebrow {
  color: #fde68a;
  font-size: 12px;
  font-weight: 900;
  letter-spacing: .17em;
  text-transform: uppercase;
}

.hero h1 {
  margin: 8px 0 7px;
  font-size: clamp(34px, 5vw, 59px);
  line-height: 1.01;
  letter-spacing: -.048em;
  color: #fffdf5;
}

.hero .gradient-word {
  background: linear-gradient(110deg, #fff7ae 0%, #facc15 42%, #f59e0b 100%);
  -webkit-background-clip: text;
  background-clip: text;
  color: transparent;
}

.hero p {
  position: relative;
  z-index: 1;
  max-width: 900px;
  color: #c8bda1;
  font-size: 16px;
  line-height: 1.65;
  margin: 0;
}

.badges {
  position: relative;
  z-index: 1;
  display: flex;
  flex-wrap: wrap;
  gap: 9px;
  margin-top: 19px;
}

.badge {
  border: 1px solid var(--line);
  background: rgba(255, 255, 255, .035);
  padding: 7px 11px;
  border-radius: 999px;
  color: #d8ccb0;
  font-size: 12px;
  backdrop-filter: blur(8px);
}

.badge strong {
  color: #fff8da;
  margin-right: 4px;
}

.app-panel {
  background: var(--surface) !important;
  border: 1px solid var(--line) !important;
  border-radius: 21px !important;
  box-shadow: 0 20px 55px rgba(0, 0, 0, .25);
  overflow: hidden;
}

.input-card {
  padding: 4px 4px 0;
}

.sidebar-card {
  background: var(--surface-2);
  border: 1px solid var(--line);
  border-radius: 18px;
  padding: 18px;
  margin-bottom: 14px;
  box-shadow: inset 0 1px 0 rgba(255, 255, 255, .02);
}

.sidebar-card h3 {
  margin: 0 0 8px;
  color: #fff8dc;
  font-size: 14px;
}

.sidebar-card p,
.sidebar-card li {
  color: var(--muted);
  font-size: 13px;
  line-height: 1.58;
}

.sidebar-card ol {
  margin: 9px 0 0;
  padding-left: 20px;
}

.model-source {
  color: #fde68a;
  overflow-wrap: anywhere;
}

#schema textarea,
#question textarea {
  font-size: 14px !important;
  line-height: 1.55 !important;
}

#sql-output {
  min-height: 285px;
}

#sql-output .cm-editor,
#sql-output textarea {
  font-size: 14px !important;
}

#generate-button {
  min-width: 155px;
  font-weight: 900;
}

button.primary,
#generate-button {
  background: linear-gradient(135deg, #eab308, #f59e0b) !important;
  color: #1b1302 !important;
  border: 1px solid rgba(255, 235, 120, .28) !important;
  box-shadow: 0 8px 24px rgba(234, 179, 8, .15) !important;
}

button.primary:hover,
#generate-button:hover {
  filter: brightness(1.07);
}

.status-card {
  border: 1px solid var(--line);
  background: rgba(255, 255, 255, .025);
  border-radius: 14px;
  padding: 12px 14px;
  color: #c9bda0;
  font-size: 12px;
  line-height: 1.55;
}

.status-card strong {
  color: #fff6cd;
}

.status-ok {
  color: var(--success);
}

.status-warn {
  color: #fde68a;
}

.status-error {
  color: var(--danger);
}

.accordion {
  background: rgba(255, 255, 255, .02) !important;
  border-color: var(--line) !important;
}

.footer-note {
  color: #8f8264;
  font-size: 11px;
  text-align: center;
  margin-top: 17px;
}

.footer-note code,
.sidebar-card code {
  color: #fde68a;
}

@media (max-width: 800px) {
  .main-shell {
    padding: 14px 10px 28px;
  }

  .hero {
    padding: 23px 20px;
    border-radius: 18px;
  }

  .hero h1 {
    font-size: 37px;
  }

  #sql-output {
    min-height: 230px;
  }
}
"""

HEAD = """
<meta name="theme-color" content="#110c03">
<meta
  name="description"
  content="Fine-tuned Qwen2.5-Coder Text-to-SQL generation with LoRA and Hugging Face."
>
"""


# Model loading and inference
def _local_adapter_available(source: str) -> bool:
    path = Path(source)
    return path.is_dir() and (path / "adapter_config.json").exists()


def _adapter_is_configured(source: str) -> bool:
    """Treat a local adapter path or non-default Hub model ID as configured."""
    if _local_adapter_available(source):
        return True

    # OUTPUT_DIR is the default local path produced by train.py. If it does not
    # exist, do not ask the Hub for a repo literally named './qwen-text-to-sql-lora'
    return source not in {"", OUTPUT_DIR, f"./{Path(OUTPUT_DIR).name}"}


def get_model():
    """Load the model once and reuse it across generations."""
    global MODEL, TOKENIZER, MODEL_MODE, MODEL_SOURCE

    if MODEL is not None and TOKENIZER is not None:
        return MODEL, TOKENIZER

    with MODEL_LOCK:
        if MODEL is not None and TOKENIZER is not None:
            return MODEL, TOKENIZER

        adapter_configured = _adapter_is_configured(ADAPTER_ID)

        # The LoRA adapter does not require a separate tokenizer vocabulary for
        # this project, so use the original Qwen tokenizer directly. This also
        # avoids depending on a duplicate large tokenizer.json inside the adapter
        print(f"[startup] Loading tokenizer: {MODEL_ID}", flush=True)
        TOKENIZER = AutoTokenizer.from_pretrained(
            MODEL_ID,
            token=HF_TOKEN,
        )
        if TOKENIZER.pad_token is None:
            TOKENIZER.pad_token = TOKENIZER.eos_token

        # ZeroGPU supports CUDA placement at module startup through CUDA
        # emulation. FP16 is sufficient for inference and avoids probing CUDA
        # capabilities before the real ZeroGPU device is attached
        if IS_HF_SPACE:
            dtype = torch.float16
            target_device = "cuda"
        elif torch.cuda.is_available():
            dtype = (
                torch.bfloat16
                if torch.cuda.is_bf16_supported()
                else torch.float16
            )
            target_device = "cuda"
        else:
            dtype = torch.float32
            target_device = "cpu"

        print(
            f"[startup] Loading base model: {MODEL_ID} "
            f"(dtype={dtype}, target_device={target_device})",
            flush=True,
        )

        base_model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            dtype=dtype,
            token=HF_TOKEN,
            low_cpu_mem_usage=True,
        )

        if adapter_configured:
            print(f"[startup] Loading LoRA adapter: {ADAPTER_ID}", flush=True)
            MODEL = PeftModel.from_pretrained(
                base_model,
                ADAPTER_ID,
                token=HF_TOKEN,
                torch_device="cpu",
            )
            MODEL_MODE = "LoRA adapter"
            MODEL_SOURCE = ADAPTER_ID
        else:
            print(
                "[startup] LoRA adapter was not found; using base-model fallback.",
                flush=True,
            )
            MODEL = base_model
            MODEL_MODE = "Base model fallback"
            MODEL_SOURCE = MODEL_ID

        MODEL = MODEL.to(target_device)
        MODEL.eval()

        print(
            f"[startup] Model ready: mode={MODEL_MODE}, source={MODEL_SOURCE}, "
            f"device={next(MODEL.parameters()).device}",
            flush=True,
        )

        return MODEL, TOKENIZER


# ZeroGPU startup model placement
if IS_HF_SPACE:
    get_model()


def _build_system_prompt(dialect_label: str) -> str:
    prompt = SYSTEM_PROMPT
    if dialect_label != "Auto / Generic SQL":
        prompt += (
            f"\n5. Generate SQL compatible with {dialect_label}."
            " Prefer syntax natural to that dialect when dialect-specific syntax is needed."
        )
    return prompt


def _clean_sql(text: str) -> str:
    """Remove common Markdown wrappers while preserving generated SQL."""
    clean = (text or "").strip()
    if clean.startswith("```"):
        clean = clean.removeprefix("```sql").removeprefix("```SQL").removeprefix("```")
        clean = clean.removesuffix("```").strip()
    return clean


def _validate_sql(sql: str, dialect_label: str) -> tuple[bool, str]:
    if not sql.strip():
        return False, "No SQL was generated."

    dialect = DIALECT_MAP.get(dialect_label)
    try:
        sqlglot.parse_one(sql, read=dialect)
        return True, "Parsed successfully with SQLGlot."
    except Exception as exc:
        return False, str(exc).split("\n", 1)[0][:220]


@spaces.GPU(duration=60)
def generate_sql_ui(
    schema: str,
    question: str,
    dialect_label: str,
    temperature: float,
    max_new_tokens: int,
):
    """Generate SQL from a schema and natural-language request."""
    clean_schema = (schema or "").strip()
    clean_question = (question or "").strip()

    if not clean_schema or not clean_question:
        missing = "database schema/context" if not clean_schema else "natural-language request"
        status = (
            "<div class='status-card status-error'>"
            f"<strong>Missing input:</strong> Please provide the {html.escape(missing)}."
            "</div>"
        )
        return "", status, {}

    model, tokenizer = get_model()

    user_message = (
        "Database context:\n"
        f"{clean_schema}\n\n"
        "Request:\n"
        f"{clean_question}"
    )

    messages = [
        {"role": "system", "content": _build_system_prompt(dialect_label)},
        {"role": "user", "content": user_message},
    ]

    prompt_text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(prompt_text, return_tensors="pt")
    device = next(model.parameters()).device
    inputs = {key: value.to(device) for key, value in inputs.items()}

    temperature = float(temperature)
    do_sample = temperature > 0.0

    generation_kwargs: dict[str, Any] = {
        "max_new_tokens": int(max_new_tokens),
        "do_sample": do_sample,
        "pad_token_id": tokenizer.eos_token_id,
        "eos_token_id": tokenizer.eos_token_id,
    }
    if do_sample:
        generation_kwargs.update(
            temperature=max(temperature, 1e-5),
            top_p=0.95,
        )

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            **generation_kwargs,
        )

    generated_tokens = outputs[0][inputs["input_ids"].shape[1] :]
    sql = _clean_sql(
        tokenizer.decode(generated_tokens, skip_special_tokens=True)
    )

    is_valid, validation_message = _validate_sql(sql, dialect_label)
    status_class = "status-ok" if is_valid else "status-warn"
    validity_text = "Valid SQL syntax" if is_valid else "Review generated SQL"

    status = f"""
    <div class="status-card">
      <strong>Generation complete</strong><br>
      <span class="{status_class}">{html.escape(validity_text)}</span>
      · {html.escape(validation_message)}
    </div>
    """

    diagnostics = {
        "model_mode": MODEL_MODE,
        "model_source": MODEL_SOURCE,
        "base_model": MODEL_ID,
        "dialect": dialect_label,
        "syntax_valid": is_valid,
        "input_tokens": int(inputs["input_ids"].shape[1]),
        "generated_tokens": int(generated_tokens.shape[0]),
        "temperature": temperature,
        "max_new_tokens": int(max_new_tokens),
        "device": str(device),
    }

    return sql, status, diagnostics


def clear_all():
    """Reset user inputs and generated outputs."""
    return (
        "",
        "",
        "",
        "<div class='status-card'>Ready for a schema and request.</div>",
        {},
    )


# Gradio app
def build_app() -> gr.Blocks:
    adapter_label = ADAPTER_ID if _adapter_is_configured(ADAPTER_ID) else "base model fallback until adapter is added"

    hero = f"""
    <div class="hero">
      <div class="eyebrow">LoRA · Transformer · Text-to-SQL</div>

      <h1>Natural language to <span class="gradient-word">SQL</span></h1>

      <p>
        Generate executable SQL from a database schema and plain-English request using
        a Qwen2.5-Coder model fine-tuned for Text-to-SQL with Hugging Face PEFT LoRA.
      </p>

      <div class="badges">
        <span class="badge"><strong>Base</strong> Qwen2.5-Coder-0.5B-Instruct</span>
        <span class="badge"><strong>Method</strong> LoRA SFT</span>
        <span class="badge"><strong>Dataset</strong> synthetic_text_to_sql</span>
        <span class="badge"><strong>Output</strong> SQL</span>
      </div>
    </div>
    """

    with gr.Blocks(title="Qwen Text-to-SQL") as demo:
        with gr.Column(elem_classes=["main-shell"]):
            gr.HTML(hero)

            with gr.Row(equal_height=False):
                with gr.Column(
                    scale=8,
                    min_width=540,
                    elem_classes=["app-panel", "input-card"],
                ):
                    schema = gr.Code(
                        label="Database schema / context",
                        language="sql",
                        value=(
                            "CREATE TABLE customers (\n"
                            "    id INTEGER PRIMARY KEY,\n"
                            "    name TEXT,\n"
                            "    country TEXT,\n"
                            "    revenue DECIMAL(12, 2)\n"
                            ");"
                        ),
                        lines=10,
                        max_lines=18,
                        interactive=True,
                        elem_id="schema",
                    )

                    question = gr.Textbox(
                        label="Natural-language request",
                        placeholder="Example: Find the five customers with the highest revenue.",
                        lines=2,
                        max_lines=5,
                        elem_id="question",
                    )

                    with gr.Row():
                        generate = gr.Button(
                            "Generate SQL",
                            variant="primary",
                            elem_id="generate-button",
                            scale=2,
                        )
                        clear = gr.Button("Clear", scale=1)

                    sql_output = gr.Code(
                        value="",
                        label="Generated SQL",
                        language="sql",
                        lines=10,
                        max_lines=20,
                        interactive=False,
                        buttons=["copy", "download"],
                        elem_id="sql-output",
                    )

                    status = gr.HTML(
                        "<div class='status-card'>Ready for a schema and request.</div>"
                    )

                    gr.Examples(
                        examples=[
                            [
                                "CREATE TABLE customers (id INTEGER, name TEXT, country TEXT, revenue DECIMAL(12,2));",
                                "Find the five customers with the highest revenue.",
                            ],
                            [
                                "CREATE TABLE orders (order_id INTEGER, customer_id INTEGER, order_date DATE, total DECIMAL(10,2));",
                                "Show monthly revenue for 2025 ordered from highest to lowest.",
                            ],
                            [
                                "CREATE TABLE employees (employee_id INTEGER, department TEXT, salary DECIMAL(10,2), hire_date DATE);",
                                "Return the average salary for each department with at least 10 employees.",
                            ],
                            [
                                "CREATE TABLE products (product_id INTEGER, category TEXT, price DECIMAL(10,2), stock INTEGER);",
                                "Find the three most expensive products in each category.",
                            ],
                        ],
                        inputs=[schema, question],
                        label="Example prompts",
                    )

                with gr.Column(scale=5, min_width=360):
                    gr.HTML(
                        f"""
                        <div class="sidebar-card">
                          <h3>Model runtime</h3>
                          <p>
                            <strong>Adapter:</strong><br>
                            <span class="model-source">{html.escape(adapter_label)}</span>
                          </p>
                        </div>

                        <div class="sidebar-card">
                          <h3>How it works</h3>
                          <ol>
                            <li>Paste the tables and columns available to the model.</li>
                            <li>Describe the query you want in natural language.</li>
                            <li>Generate SQL and inspect the syntax validation result.</li>
                          </ol>
                        </div>
                        """
                    )

                    with gr.Accordion(
                        "Generation controls",
                        open=True,
                        elem_classes=["accordion"],
                    ):
                        dialect = gr.Dropdown(
                            choices=list(DIALECT_MAP.keys()),
                            value="Auto / Generic SQL",
                            label="SQL dialect",
                        )

                        temperature = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            value=0.0,
                            step=0.05,
                            label="Temperature",
                            info="0 is deterministic and recommended for SQL generation.",
                        )

                        max_new_tokens = gr.Slider(
                            minimum=64,
                            maximum=768,
                            value=256,
                            step=32,
                            label="Maximum output tokens",
                        )

                    with gr.Accordion(
                        "Diagnostics",
                        open=False,
                        elem_classes=["accordion"],
                    ):
                        diagnostics = gr.JSON(
                            value={},
                            label="Generation diagnostics",
                        )

                    gr.HTML(
                        """
                        <div class="sidebar-card">
                          <h3>Validation</h3>
                          <p>
                            The generated query is parsed with SQLGlot for syntax validation.
                            Syntax validity does not guarantee that the query is logically correct
                            for your database or returns the intended rows.
                          </p>
                        </div>
                        """
                    )

            gr.HTML(
                """
                <div class="footer-note">
                  Qwen2.5-Coder · Hugging Face Transformers · TRL · PEFT LoRA · Gradio
                </div>
                """
            )

        generation_inputs = [
            schema,
            question,
            dialect,
            temperature,
            max_new_tokens,
        ]
        generation_outputs = [
            sql_output,
            status,
            diagnostics,
        ]

        generate.click(
            fn=generate_sql_ui,
            inputs=generation_inputs,
            outputs=generation_outputs,
        )

        question.submit(
            fn=generate_sql_ui,
            inputs=generation_inputs,
            outputs=generation_outputs,
        )

        clear.click(
            fn=clear_all,
            outputs=[schema, question, sql_output, status, diagnostics],
        )

    return demo


if __name__ == "__main__":
    app = build_app()
    app.queue(default_concurrency_limit=1).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        show_error=True,
        ssr_mode=False,
        theme=gr.themes.Base(),
        css=CSS,
        head=HEAD,
    )