"""
SQL Agent Tools — Execute SQL queries and calculator operations.

Extracted from api/ai_sql_trace_route.py for clean separation.
"""

import asyncio
import re
from typing import Any

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

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

_FORBIDDEN = 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 = True) -> 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"
    if _FORBIDDEN.search(sql_clean):
        return "Forbidden SQL keyword"
    check = sql_clean.lower().replace("_text_embedding", "").replace("_embedding", "")
    if re.search(r"\b(embedding|vector)\b", check):
        return "Cannot query embedding/vector column"
    if require_limit and "limit" not in sql_clean.lower():
        return "Query must include LIMIT"
    return None


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

def serialize_row(row: dict[str, Any]) -> dict[str, Any]:
    """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[str, Any]) -> 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)
        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": []}

    if tool_name == "sql_postgres":
        sql = params.get("sql", "")
        err = validate_sql(sql)
        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": []}

    if 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}
            return {"error": "Invalid expression", "expression": expr}
        except Exception as e:
            return {"error": str(e), "expression": expr}

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


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


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

def make_result_preview(result: dict[str, Any] | Exception) -> dict[str, Any]:
    """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 "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)
