"""The hybrid translation path, shared by `orashift translate` and the demos.

Rules first, model second. That ordering is the strategy that scored best in
evaluation (87% against the fine-tuned model's 86%), because `sqlglot` is free,
instant and deterministic, and the model only has to handle what rules cannot.

This module depends only on the standard library and `sqlglot` (the model path
imports torch lazily), so the static demo can run it in the browser under Pyodide.

It also makes the demo usable on a free CPU. Most statements never reach the
model at all, so they return immediately; only the Oracle-specific ones pay for
a generation.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from functools import lru_cache
from typing import Any

import sqlglot
from sqlglot.errors import ParseError, TokenError, UnsupportedError

logging.getLogger("sqlglot").setLevel(logging.ERROR)

BASE_MODEL = os.environ.get("DEMO_BASE_MODEL", "Qwen/Qwen2.5-Coder-3B-Instruct")
ADAPTER = os.environ.get("DEMO_ADAPTER", "")
# bfloat16 halves memory against float32 (about 6 GB instead of 12 for a 3B model),
# which is what lets the model fit comfortably on a free 16 GB CPU Space.
DTYPE = os.environ.get("DEMO_DTYPE", "bfloat16")
MAX_NEW_TOKENS = 448

SYSTEM_PROMPT = """You translate Oracle SQL into PostgreSQL.

Rules:
- Return only the PostgreSQL statement. No explanation, no markdown, no semicolon.
- The translation must run on PostgreSQL 17 and return exactly what the Oracle
  statement returns on Oracle.
- Oracle types map as: NUMBER(p,0) to smallint/integer/bigint by precision,
  NUMBER(p,s) to numeric(p,s), VARCHAR2(n) to varchar(n), CHAR(n) to char(n),
  CLOB to text, DATE to timestamp(0).
- Watch the semantics, not just the syntax: an Oracle DATE carries a time,
  '' is NULL in Oracle, concatenating NULL yields the other operand in Oracle
  but NULL in PostgreSQL, and integer division truncates in PostgreSQL."""

# Which constructs get routed to the model.
#
# Derived from results/metrics.json, not from intuition. A first version of this
# list was written from memory and sent NVL and DECODE to the model, which was
# wrong: sqlglot scores 100% on both. The measured per-construct accuracy is:
#
#   routed to the model (sqlglot <= 83%):
#     rownum 0%, connect_by 0%, trunc_date 0%, add_months 0%, pivot 0%,
#     merge 0%, dual 0%, sequence 0%, DDL table/constraint/index 0%,
#     sysdate 8%, listagg 60%, dml_update 67%, outer_join_plus 71%,
#     null_concat 83%
#
#   left to the rules (sqlglot >= 86%):
#     nvl 100%, decode 100%, to_char_date 100%, string_func 100%,
#     analytic 100%, subquery 100%, cte 100%, aggregate 100%,
#     fetch_first 98%, minus 86%
#
# The threshold is deliberately cautious. The demo cannot execute anything, so
# it cannot tell a good rule output from a bad one; where the rules are only
# usually right, a slower answer beats a confidently wrong one.
ORACLE_ONLY = (
    "rownum",
    "connect by",
    "start with",
    "sys_connect_by_path",
    "connect_by_root",
    "connect_by_isleaf",
    "trunc(",
    "add_months",
    "months_between",
    "last_day",
    "sysdate",
    "systimestamp",
    "pivot",
    "unpivot",
    "merge into",
    "from dual",
    ".nextval",
    ".currval",
    "listagg",
    "(+)",
    "||",
    "update ",
    "varchar2",
    "number(",
    "clob",
)


@dataclass(frozen=True, slots=True)
class Translation:
    postgres_sql: str
    path: str
    """Which route produced it: 'sqlglot', 'fine-tuned model', or 'failed'."""
    note: str = ""


def needs_model(oracle_sql: str) -> list[str]:
    """Oracle-only constructs present in the statement, if any."""
    lowered = f" {' '.join(oracle_sql.lower().split())} "
    return [token.strip() for token in ORACLE_ONLY if token in lowered]


def try_sqlglot(oracle_sql: str) -> str | None:
    try:
        converted = sqlglot.transpile(oracle_sql, read="oracle", write="postgres")
    except (ParseError, TokenError, UnsupportedError, RecursionError, ValueError, KeyError):
        return None
    if not converted or not converted[0].strip():
        return None
    return converted[0].strip()


@lru_cache(maxsize=1)
def load_model() -> tuple[Any, Any] | None:
    """Load the fine-tuned model once, or return None if it is unavailable.

    The demo must still work without it: the rule path handles most statements,
    and a demo that refuses to start because a 6 GB download failed is worse
    than one that is honest about which half is running.
    """
    if not ADAPTER:
        return None
    try:
        import torch
        from peft import PeftModel
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except ImportError:
        return None

    try:
        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
        model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL, torch_dtype=getattr(torch, DTYPE), device_map="cpu"
        )
        model = PeftModel.from_pretrained(model, ADAPTER)
        model.eval()
        return model, tokenizer
    except Exception:
        return None


def preload() -> None:
    """Start loading the model in the background, so the first request does not
    wait for a multi-gigabyte download. Does nothing if no adapter is set."""
    if ADAPTER:
        import threading

        threading.Thread(target=load_model, daemon=True).start()


def run_model(oracle_sql: str, schema_context: str = "") -> str | None:
    loaded = load_model()
    if loaded is None:
        return None
    model, tokenizer = loaded

    system = SYSTEM_PROMPT
    if schema_context.strip():
        system = f"{system}\n\nOracle table definitions:\n\n{schema_context.strip()}"

    prompt = tokenizer.apply_chat_template(
        [{"role": "system", "content": system}, {"role": "user", "content": oracle_sql.strip()}],
        tokenize=False,
        add_generation_prompt=True,
    )
    import torch

    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )
    generated = output[0][inputs["input_ids"].shape[1] :]
    return tokenizer.decode(generated, skip_special_tokens=True).strip()


def translate(oracle_sql: str, schema_context: str = "") -> Translation:
    """Rules first, model as fallback. Returns the result and which path produced it."""
    oracle_sql = (oracle_sql or "").strip().rstrip(";")
    if not oracle_sql:
        return Translation("", "failed", "Nothing to translate.")

    oracle_isms = needs_model(oracle_sql)
    rule_output = try_sqlglot(oracle_sql)

    if rule_output and not oracle_isms:
        return Translation(rule_output, "sqlglot", "No Oracle-specific constructs found.")

    model_output = run_model(oracle_sql, schema_context)
    if model_output:
        found = ", ".join(f"`{t}`" for t in oracle_isms[:4]) or "an unsupported construct"
        return Translation(model_output, "fine-tuned model", f"Rules cannot handle {found}.")

    if rule_output:
        found = ", ".join(f"`{t}`" for t in oracle_isms[:4])
        return Translation(
            rule_output,
            "sqlglot (unreliable here)",
            f"The model is not loaded, and this statement contains {found}, "
            "which sqlglot usually passes through unchanged. Verify before using.",
        )

    return Translation("", "failed", "Neither path could translate this statement.")