from __future__ import annotations

import json
import re
import sys
from pathlib import Path

import sqlparse

from felderize.config import Config
from felderize.constants import DEFAULT_MAX_RETRIES
from felderize.feldera_client import validate_sql
from felderize.llm import LLMClient, PromptTooLargeError, create_client
from felderize.models import Status, TranslationResult
from felderize.prompts import (
    build_force_query_prompt,
    build_repair_prompt,
    build_user_prompt,
)
from felderize.skills import build_system_prompt


_CREATE_TABLE_RE = re.compile(r"\bCREATE\s+TABLE\b", re.IGNORECASE)
_CREATE_VIEW_RE = re.compile(
    r"\bCREATE\s+(OR\s+REPLACE\s+)?(TEMP(ORARY)?\s+)?VIEW\b", re.IGNORECASE
)
_SELECT_RE = re.compile(r"\bSELECT\b", re.IGNORECASE)


def trim_schema_to_referenced(schema_sql: str, query_sql: str) -> str:
    """Return schema_sql with unreferenced CREATE TABLE statements removed.

    Builds a reference set from the query SQL and the SELECT bodies of any
    CREATE VIEW statements in the schema. CREATE TABLE definitions themselves
    are excluded from the scan so a table's own name doesn't count as a
    self-reference. Non-table statements (views) always pass through.

    This reduces LLM prompt size when a schema has many shared tables that the
    specific query doesn't use.
    """
    stmts = [s.strip() for s in sqlparse.split(schema_sql) if s.strip()]

    scan_parts: list[str] = [query_sql]
    for stmt in stmts:
        s = stmt.rstrip(";")
        vm = re.match(
            r"CREATE\s+(?:OR\s+REPLACE\s+)?(?:TEMP(?:ORARY)?\s+)?VIEW\s+\S+"
            r"(?:\s*\([^)]*\))?"  # optional column list
            r"(?:\s|--[^\n]*\n)*"  # optional whitespace and line comments
            r"AS\s+(.*)",
            s,
            re.IGNORECASE | re.DOTALL,
        )
        if vm:
            scan_parts.append(vm.group(1))

    ref_tokens: set[str] = set(
        re.findall(r"\b[a-z_]\w*\b", "\n".join(scan_parts).lower())
    )

    kept: list[str] = []
    for stmt in stmts:
        s = stmt.rstrip(";")
        m = re.match(
            r"CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?(\w+)", s, re.IGNORECASE
        )
        if m:
            if m.group(1).lower() in ref_tokens:
                kept.append(stmt if stmt.endswith(";") else stmt + ";")
        else:
            kept.append(stmt if stmt.endswith(";") else stmt + ";")

    return "\n\n".join(kept)


def validate_schema(schema_sql: str) -> list[str]:
    """Return validation errors for the schema SQL, empty if valid."""
    if not schema_sql.strip():
        return ["Schema file is empty."]
    if not _CREATE_TABLE_RE.search(schema_sql):
        return [
            "Schema file does not contain any CREATE TABLE statement. "
            "Expected a SQL schema, not arbitrary text."
        ]
    return []


def validate_query(query_sql: str) -> list[str]:
    """Return validation errors for the query SQL, empty if valid."""
    if not query_sql.strip():
        return ["Query file is empty."]
    if not _CREATE_VIEW_RE.search(query_sql) and not _SELECT_RE.search(query_sql):
        return [
            "Query file does not contain any CREATE VIEW or SELECT statement. "
            "Expected a SQL query, not arbitrary text."
        ]
    return []


def validate_inputs(schema_sql: str, query_sql: str) -> list[str]:
    """Return validation errors for both schema and query SQL."""
    return validate_schema(schema_sql) + validate_query(query_sql)


def split_combined_sql(sql: str) -> tuple[str, str]:
    """Split a combined SQL file into (schema_sql, query_sql).

    Schema: CREATE TABLE statements. Query: CREATE [OR REPLACE] [TEMP] VIEW
    statements. Unrecognised statements are placed in schema_sql.
    """
    raw_stmts = [s.strip().rstrip(";") for s in sqlparse.split(sql) if s.strip()]

    schema_parts: list[str] = []
    query_parts: list[str] = []

    for stripped in raw_stmts:
        if not stripped:
            continue
        first_kw = next(
            (
                ln.strip()
                for ln in stripped.splitlines()
                if ln.strip() and not ln.strip().startswith("--")
            ),
            "",
        ).upper()
        if not first_kw:
            continue
        if _CREATE_VIEW_RE.match(first_kw):
            query_parts.append(stripped + ";")
        else:
            schema_parts.append(stripped + ";")

    return "\n\n".join(schema_parts), "\n\n".join(query_parts)


# ---------------------------------------------------------------------------
# LLM response parsing
# ---------------------------------------------------------------------------


def _parse_response(raw: str) -> dict:
    """Extract JSON from an LLM response, handling markdown fences."""
    text = raw.strip()
    match = re.search(r"```(?:json)?\s*\n(.*?)```", text, re.DOTALL)
    if match:
        text = match.group(1).strip()
    return json.loads(text)


_SQL_FENCE_RE = re.compile(r"```(?:sql)?\s*\n(.*?)```", re.DOTALL | re.IGNORECASE)

# A salvageable statement must be a *structured* CREATE — a VIEW followed by AS,
# or a TABLE followed by its column list — not merely the word "CREATE". The
# strict shape stops the model's reasoning prose (e.g. "CREATE VIEW statement
# from the previous attempt's reasoning...") from being scraped in as SQL.
_CREATE_STMT_RE = re.compile(
    r"""\bCREATE\s+(?:OR\s+REPLACE\s+)?
        (?:(?:LOCAL|MATERIALIZED|TEMP(?:ORARY)?)\s+)?
        (?:VIEW\s+[\w`".-]+\s+AS\b | TABLE\s+[\w`".-]+\s*\()
        .*?;""",
    re.DOTALL | re.IGNORECASE | re.VERBOSE,
)


def _salvage_sql(raw: str) -> tuple[str, str]:
    """Best-effort extraction of (schema_sql, query_sql) from a non-JSON response.

    Some responses come back as prose or bare SQL instead of the requested JSON
    (or as truncated JSON). Rather than give up and return nothing, pull whatever
    SQL we can: prefer ```sql fences, else scan for well-formed CREATE statements,
    then split them into schema (CREATE TABLE) vs query (CREATE VIEW). Returns
    ("", "") when no SQL-looking statements are present.
    """
    fences = _SQL_FENCE_RE.findall(raw)
    text = "\n".join(f.strip() for f in fences) if fences else raw
    statements = _CREATE_STMT_RE.findall(text)
    if not statements:
        return "", ""
    return split_combined_sql("\n".join(statements))


def _parse_or_salvage(raw: str) -> tuple[dict, list[str]]:
    """Parse the JSON response, or salvage best-effort SQL when that fails.

    Returns (data, warnings). On a clean parse, warnings is empty. On failure,
    salvaged SQL (if any) is returned as a result dict so the caller always has
    something to validate/repair rather than an empty translation.
    """
    try:
        return _parse_response(raw), []
    except (json.JSONDecodeError, KeyError):
        schema_sql, query_sql = _salvage_sql(raw)
        if schema_sql.strip() or query_sql.strip():
            return (
                {
                    "feldera_schema": schema_sql,
                    "feldera_query": query_sql,
                    "unsupported": [],
                },
                ["LLM response was not valid JSON — salvaged SQL from the raw text"],
            )
        return {}, ["Failed to parse LLM response and no SQL could be salvaged"]


def _as_str(val: object) -> str:
    """Normalize a value to string — handles lists from LLM responses."""
    if isinstance(val, list):
        return "\n".join(str(v) for v in val)
    return str(val) if val else ""


def _as_list(val: object) -> list[str]:
    if isinstance(val, list):
        return [str(v) for v in val]
    if isinstance(val, str) and val:
        return [val]
    return []


def _normalize_escapes(sql: str) -> str:
    """Undo a doubly-escaped LLM response.

    Some responses arrive escaped a second time, so the SQL body holds the
    literal two-character sequences ``\\n`` and ``\\"`` instead of real newlines
    and quotes; the compiler then fails on the stray backslash. When the text
    carries no real newline yet contains a literal ``\\n``, the whole statement
    has collapsed onto one line — treat it as doubly escaped and restore the
    intended characters. Genuinely multi-line SQL is left untouched.
    """
    if "\n" in sql or "\\n" not in sql:
        return sql
    return (
        sql.replace("\\n", "\n")
        .replace("\\t", "\t")
        .replace("\\r", "\r")
        .replace('\\"', '"')
    )


def _build_result(data: dict) -> TranslationResult:
    unsupported = _as_list(data.get("unsupported", []))
    return TranslationResult(
        feldera_schema=_normalize_escapes(_as_str(data.get("feldera_schema", ""))),
        feldera_query=_normalize_escapes(_as_str(data.get("feldera_query", ""))),
        unsupported=unsupported,
        warnings=_as_list(data.get("warnings", [])),
        explanations=_as_list(data.get("explanations", [])),
        status=Status.UNSUPPORTED if unsupported else Status.SUCCESS,
    )


def _too_large_message(detail: str) -> str:
    """User-facing guidance for a program that overflows the model's context window."""
    return (
        "The program is too large to fit the model's context window "
        f"({detail}). Shorten the input — translate fewer views at a time, drop "
        "unused tables, or split it into smaller files — or set a model with a "
        "larger context window via FELDERIZE_MODEL / --model."
    )


def _log_verbose_sql(full_sql: str, attempt: int) -> None:
    print(f"\n--- SQL submitted to validator (attempt {attempt}) ---", file=sys.stderr)
    print(full_sql, file=sys.stderr)
    print("---", file=sys.stderr)


# ---------------------------------------------------------------------------
# Translation: one LLM call on the whole program, with optional compiler repair
# ---------------------------------------------------------------------------


def _translate_with_repair(
    schema_sql: str,
    query_sql: str,
    config: Config,
    client: LLMClient,
    system_prompt: str,
    validate: bool,
    max_retries: int,
    verbose: bool = False,
) -> TranslationResult:
    """Translate the whole program with one LLM call, then optionally validate + repair."""
    raw = client.translate(system_prompt, build_user_prompt(schema_sql, query_sql))
    data, parse_warnings = _parse_or_salvage(raw)
    result = _build_result(data)
    result.warnings.extend(parse_warnings)

    # If the query came back empty, force one retry with NULL placeholders.
    if not result.feldera_query.strip():
        raw2 = client.translate(
            system_prompt,
            build_force_query_prompt(schema_sql, query_sql, result.unsupported),
        )
        data2, warn2 = _parse_or_salvage(raw2)
        result2 = _build_result(data2)
        if result2.feldera_query.strip():
            result = result2
            result.warnings.extend(warn2)
            result.warnings.append(
                "Query produced on force-retry after empty first response"
            )

    # A response we could not parse or salvage is a genuine failure, not a clean
    # translation — mark it so callers don't treat an empty query as success.
    if parse_warnings and not result.feldera_query.strip():
        result.status = Status.ERROR

    if not validate:
        return result

    full_sql = result.feldera_schema + "\n\n" + result.feldera_query
    # Schema-only or empty SQL always compiles, which is not a meaningful success.
    if not full_sql.strip() or not result.feldera_query.strip():
        if not result.unsupported:
            result.status = Status.UNSUPPORTED
            result.unsupported = ["No query generated — translation incomplete"]
        return result

    for attempt in range(max_retries):
        if verbose:
            _log_verbose_sql(full_sql, attempt + 1)
        errors = validate_sql(full_sql, config.compiler_path)
        if errors and any("Compiler not found" in e for e in errors):
            print(
                "Warning: Feldera compiler not found — skipping validation.",
                file=sys.stderr,
            )
            result.warnings.append("Compiler not found — output SQL is not validated")
            result.status = Status.UNSUPPORTED if result.unsupported else Status.SUCCESS
            return result
        if not errors:
            result.warnings.append(f"Validated successfully (attempt {attempt + 1})")
            result.status = Status.UNSUPPORTED if result.unsupported else Status.SUCCESS
            return result

        print(
            f"Validation attempt {attempt + 1}/{max_retries} failed: "
            f"{len(errors)} error(s)",
            file=sys.stderr,
        )
        for err in errors:
            print(f"  {err}", file=sys.stderr)

        repair = build_repair_prompt(schema_sql, query_sql, full_sql, errors)
        raw = client.translate(system_prompt, repair)
        data, repair_warnings = _parse_or_salvage(raw)
        if repair_warnings:
            result.warnings.append(
                f"Repair attempt {attempt + 1}: {repair_warnings[0]}"
            )
        # When parsing/salvage yields nothing, data is empty and these .get calls
        # keep the previous best-effort SQL rather than blanking it out.
        result.feldera_schema = _as_str(
            data.get("feldera_schema", result.feldera_schema)
        )
        result.feldera_query = _as_str(data.get("feldera_query", result.feldera_query))
        result.unsupported = _as_list(data.get("unsupported", result.unsupported))
        result.explanations = _as_list(data.get("explanations", result.explanations))
        full_sql = result.feldera_schema + "\n\n" + result.feldera_query

    if verbose:
        _log_verbose_sql(full_sql, max_retries + 1)
    errors = validate_sql(full_sql, config.compiler_path)
    if not errors:
        result.warnings.append(f"Validated successfully (attempt {max_retries + 1})")
        result.status = Status.UNSUPPORTED if result.unsupported else Status.SUCCESS
    else:
        result.status = Status.ERROR
        result.warnings.extend(
            f"Still failing after {max_retries} repairs: {e}" for e in errors
        )

    return result


def translate_spark_to_feldera(
    schema_sql: str,
    query_sql: str,
    config: Config,
    validate: bool = False,
    max_retries: int = DEFAULT_MAX_RETRIES,
    include_docs: bool = True,
    verbose: bool = False,
    extra_rules: list[Path] | None = None,
    extra_examples_dirs: list[Path] | None = None,
    extra_examples_files: list[Path] | None = None,
) -> TranslationResult:
    """Translate a Spark SQL schema + query (the whole program) to Feldera SQL.

    A single LLM call translates the entire program. When ``validate`` is set the
    output is checked against the Feldera compiler and repaired with error
    feedback for up to ``max_retries`` attempts. On a non-success first pass the
    translation is retried once more with Feldera documentation added to the
    prompt (unless ``include_docs`` is false).
    """
    schema_sql = trim_schema_to_referenced(schema_sql, query_sql)
    combined_sql = schema_sql + "\n" + query_sql
    client = create_client(config, verbose=verbose)

    # First pass: rules + validated examples (no docs).
    system_prompt = build_system_prompt(
        spark_sql=combined_sql,
        with_docs=False,
        with_examples=True,
        extra_rules=extra_rules,
        extra_examples_dirs=extra_examples_dirs,
        extra_examples_files=extra_examples_files,
    )
    try:
        result = _translate_with_repair(
            schema_sql,
            query_sql,
            config,
            client,
            system_prompt,
            validate,
            max_retries,
            verbose,
        )
    except PromptTooLargeError as e:
        message = _too_large_message(str(e))
        print(f"Error: {message}", file=sys.stderr)
        return TranslationResult(status=Status.ERROR, warnings=[message])

    # Retry with docs when the first pass was not fully successful.
    if not include_docs or result.status == Status.SUCCESS:
        return result

    print("Retrying with skills + docs prompt...", file=sys.stderr)
    system_prompt_docs = build_system_prompt(
        spark_sql=combined_sql,
        with_docs=True,
        with_examples=True,
        with_skills=True,
        retry_unsupported=result.unsupported or None,
        extra_rules=extra_rules,
        extra_examples_dirs=extra_examples_dirs,
        extra_examples_files=extra_examples_files,
    )
    try:
        docs_result = _translate_with_repair(
            schema_sql,
            query_sql,
            config,
            client,
            system_prompt_docs,
            validate,
            max_retries,
            verbose,
        )
    except PromptTooLargeError as e:
        # The docs prompt is the largest; keep the first-pass result and explain
        # why the docs retry was skipped.
        message = _too_large_message(str(e))
        print(f"Warning: docs retry skipped — {message}", file=sys.stderr)
        result.warnings.append(f"Docs retry skipped — {message}")
        return result

    result = docs_result
    if result.status != Status.ERROR:
        result.warnings.append("Resolved with skills + docs retry")
    return result