"""
SQL Agent Core — Shared logic for Text-to-SQL, Dashboard, and SQL Trace agents.

Extracted from api/sql_chat_route.py and api/ai_sql_trace_route.py.
Contains: SQL validation, LLM calls, tool execution, result serialization.
"""

import asyncio
import json
import logging
import re
from typing import Any

import httpx

from common.starrocks_connection import StarRocksConnection
from common.postgres_readonly import PostgresReadonly

logger = logging.getLogger(__name__)


# ─── SQL Safety ──────────────────────────────────────────────────────

_FORBIDDEN_PATTERNS = re.compile(
    r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|REPLACE|MERGE|GRANT|REVOKE)\b",
    re.IGNORECASE,
)


def validate_sql(sql: str, *, require_limit: bool = False) -> str | None:
    """Validate SQL is safe (SELECT only). Returns error message or None."""
    sql_clean = sql.strip().rstrip(";")
    upper_sql = sql_clean.upper()
    if not any(upper_sql.startswith(k) for k in ("SELECT", "WITH", "SHOW")):
        return "Only SELECT, WITH, or SHOW queries allowed"
    match = _FORBIDDEN_PATTERNS.search(sql_clean)
    if match:
        return f"Forbidden SQL keyword: {match.group()}"
    check_str = sql_clean.lower().replace("_text_embedding", "").replace("_embedding", "")
    if re.search(r'\b(embedding|vector)\b', check_str):
        return "Cannot query embedding/vector column"
    if require_limit and "limit" not in sql_clean.lower():
        return "Query must include LIMIT"
    return None


# ─── LLM Call (Codex) ───────────────────────────────────────────────

async def call_llm(
    system_prompt: str,
    user_content: str,
    model_id: str = "codex/gpt-5.3-codex",
) -> str:
    """Call Codex LLM and return raw text response."""
    from common.codex_auth import get_codex_access_token, get_codex_account_id, get_refresh_token
    from common.codex_chat_model import CodexChatModel
    from langchain_core.messages import HumanMessage, SystemMessage

    access_token = get_codex_access_token()
    if not access_token:
        raise RuntimeError("Codex auth not available. Run `codex login` first.")

    actual_model = model_id.replace("codex/", "") if model_id.startswith("codex/") else model_id
    logger.info("🔧 SQL Agent model: %s", actual_model)

    model = CodexChatModel(
        model_name=actual_model,
        access_token=access_token,
        account_id=get_codex_account_id() or "",
        refresh_token=get_refresh_token(),
        streaming=False,
    )

    messages = [SystemMessage(content=system_prompt), HumanMessage(content=user_content)]
    result = await model._agenerate(messages)
    return result.generations[0].message.content


# ─── JSON Parsing ────────────────────────────────────────────────────

def parse_json(raw: str) -> dict:
    """Parse JSON from LLM response, stripping markdown fences."""
    json_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', raw)
    json_str = json_match.group(1).strip() if json_match else raw.strip()
    return json.loads(json_str)


# ─── Row Serialization ──────────────────────────────────────────────

def serialize_row(row: dict) -> dict:
    """Serialize a row for JSON output (handle Decimal, bytes, datetime, NaN)."""
    for key, val in row.items():
        if isinstance(val, (bytes, bytearray)):
            row[key] = val.hex()
        elif hasattr(val, 'isoformat'):
            row[key] = val.isoformat()
        elif isinstance(val, (float,)) and (val != val):
            row[key] = 0
        elif hasattr(val, 'as_tuple'):  # Decimal
            row[key] = float(val)
    return row


# ─── Tool Executors ──────────────────────────────────────────────────

async def exec_tool(tool_name: str, params: dict, *, require_limit: bool = False) -> dict[str, Any]:
    """Execute a single tool and return results."""

    if tool_name in ("sql_langfuse", "sql_starrocks"):
        sql = params.get("sql", "")
        err = validate_sql(sql, require_limit=require_limit)
        if err:
            return {"error": err, "data": []}
        db = StarRocksConnection()
        try:
            rows = await db.execute_query_async(sql)
            data = [serialize_row(dict(r)) for r in rows]
            return {"data": data, "columns": list(rows[0].keys()) if rows else [], "row_count": len(data)}
        except Exception as e:
            return {"error": str(e)[:300], "data": []}

    elif tool_name == "sql_postgres":
        sql = params.get("sql", "")
        err = validate_sql(sql, require_limit=require_limit)
        if err:
            return {"error": err, "data": []}
        try:
            rows = await PostgresReadonly.execute_query_async(sql)
            data = [serialize_row(dict(r)) for r in rows]
            return {"data": data, "columns": list(rows[0].keys()) if rows else [], "row_count": len(data)}
        except Exception as e:
            return {"error": str(e)[:300], "data": []}

    elif tool_name == "calculator":
        expr = params.get("expression", "")
        description = params.get("description", "")
        try:
            allowed = set("0123456789.+-*/() ,")
            if all(c in allowed for c in expr.replace(" ", "")):
                result = eval(expr)  # noqa: S307
                return {"expression": expr, "result": result, "description": description}
            else:
                return {"error": "Invalid expression (only math allowed)", "expression": expr}
        except Exception as e:
            return {"error": str(e), "expression": expr}

    else:
        return {"error": f"Unknown tool: {tool_name}", "data": []}


async def execute_tools_parallel(tools: list[dict], *, require_limit: bool = False) -> list[dict]:
    """Execute multiple tools in parallel, return results."""
    tasks = [exec_tool(t.get("name", ""), t.get("params", {}), require_limit=require_limit) for t in tools]
    return await asyncio.gather(*tasks, return_exceptions=True)


# ─── Result Preview & Summary ────────────────────────────────────────

def make_result_preview(result: dict | Exception) -> dict:
    """Create a compact preview of a tool result for SSE streaming."""
    preview: dict[str, Any] = {}
    if isinstance(result, Exception):
        return {"error": str(result)[:200]}
    if "data" in result:
        data = result["data"]
        if isinstance(data, list):
            preview["row_count"] = len(data)
            preview["columns"] = result.get("columns", [])
            preview["preview_rows"] = data[:5] if data else []
    if "summary" in result:
        preview["summary"] = result["summary"]
    if "error" in result:
        preview["error"] = result["error"]
    if "result" in result:
        preview["result"] = result["result"]
    return preview


def summarize_results(all_results: dict[str, Any]) -> str:
    """Compress tool results into concise format optimized for token efficiency."""
    parts = []
    for key, result in all_results.items():
        if isinstance(result, Exception):
            parts.append(f"\n## {key}\nERROR: {str(result)[:200]}")
            continue

        result_dict = dict(result) if isinstance(result, dict) else {"error": str(result)}

        if "error" in result_dict:
            parts.append(f"\n## {key}\nERROR: {result_dict['error']}")
            continue

        if "result" in result_dict:
            parts.append(f"\n## {key}\nResult: {result_dict['result']}")
            continue

        data = result_dict.get("data", [])
        if not isinstance(data, list) or not data:
            parts.append(f"\n## {key}\n(no data)")
            continue

        columns = result_dict.get("columns", list(data[0].keys()) if data else [])
        row_count = len(data)

        if row_count <= 20:
            lines = [f"\n## {key} ({row_count} rows)"]
            lines.append("\t".join(str(c) for c in columns))
            for row in data:
                cells = []
                for col in columns:
                    val = row.get(col, "")
                    s = str(val) if val is not None else ""
                    if len(s) > 60:
                        s = s[:57] + "..."
                    cells.append(s)
                lines.append("\t".join(cells))
            parts.append("\n".join(lines))
        else:
            lines = [f"\n## {key} ({row_count} rows → auto-summarized)"]
            lines.append("### Sample (first 10):")
            lines.append("\t".join(str(c) for c in columns))
            for row in data[:10]:
                cells = []
                for col in columns:
                    val = row.get(col, "")
                    s = str(val) if val is not None else ""
                    if len(s) > 40:
                        s = s[:37] + "..."
                    cells.append(s)
                lines.append("\t".join(cells))

            # Compute stats for numeric columns
            lines.append("### Statistics:")
            for col in columns:
                values = [row.get(col) for row in data if row.get(col) is not None]
                if not values:
                    continue
                try:
                    nums = [float(v) for v in values if v != "" and v is not None]
                    if nums and len(nums) > len(values) * 0.5:
                        lines.append(
                            f"  {col}: min={min(nums):.2f}, max={max(nums):.2f}, "
                            f"avg={sum(nums)/len(nums):.2f}, sum={sum(nums):.2f}, count={len(nums)}"
                        )
                except (ValueError, TypeError):
                    unique = list(set(str(v) for v in values[:100]))
                    top = unique[:5]
                    lines.append(f"  {col}: {len(unique)} unique values, top: {', '.join(top)}")

            parts.append("\n".join(lines))

    return "\n".join(parts)


# ─── Error Humanizer ─────────────────────────────────────────────────

def humanize_llm_error(exc: Exception) -> str:
    """Convert LLM exceptions to human-readable Vietnamese messages."""
    if isinstance(exc, httpx.HTTPStatusError):
        response = exc.response
        status = response.status_code if response else "unknown"
        body = ""
        if response is not None:
            try:
                body = response.text[:300]
            except Exception:
                body = str(exc)

        lowered = body.lower()
        if "no route to host" in lowered:
            return (
                f"Codex backend is unreachable right now (HTTP {status}: No route to host). "
                "The service cannot connect to the upstream model."
            )
        if status == 503:
            return "Codex backend is temporarily unavailable (HTTP 503). Please retry in a moment."
        return f"Codex API error (HTTP {status}). {body[:180]}".strip()

    if isinstance(exc, httpx.RequestError):
        return f"Network error while calling Codex backend: {exc}"

    return str(exc)
