"""
Report Agent Core — Shared utilities for the HTML Report Agent.

Contains: LLM calls, SQL validation, tool execution, result summarization,
temp report storage, and SSE formatting.

Used by: agent/report_agent/graph.py, api/report_html_route.py
"""

import asyncio
import json
import logging
import re
from datetime import datetime, timezone
from itertools import count
from typing import Any, AsyncGenerator

import httpx

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

logger = logging.getLogger(__name__)


# ─── 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) -> 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"
    return None


# ─── LLM Call (Codex + OpenAI) ──────────────────────────────────────

async def call_llm(
    system_prompt: str,
    user_content: str,
    model_id: str = "codex/gpt-5.3-codex",
    codex_token: str | None = None,
    openai_key: str | None = None,
    json_mode: bool = False,
) -> str:
    """Call LLM and return raw text response. Supports codex/* and openai/* models.

    Args:
        json_mode: If True, enforce JSON output format via OpenAI response_format
                   or Gemini response_mime_type. Codex models rely on prompt enforcement.
    """
    if model_id.startswith("openai/"):
        actual_model = model_id.replace("openai/", "")
        from common.llm_factory import LLMFactory

        if not openai_key:
            raise RuntimeError("OpenAI API key không tồn tại. Thêm Token trong Cài Đặt (Sidebar).")

        factory = LLMFactory()
        model = factory.get_model(
            model_name=actual_model, streaming=False,
            json_mode=json_mode, api_key=openai_key,
        )
        from langchain_core.messages import HumanMessage, SystemMessage

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

    else:
        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 = codex_token or get_codex_access_token()
        if not access_token:
            raise RuntimeError("Codex auth not available. Thêm Token trong Cài Đặt (Sidebar).")

        actual_model = model_id.replace("codex/", "") if model_id.startswith("codex/") else model_id

        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


async def call_llm_streaming(
    system_prompt: str,
    user_content: str,
    model_id: str = "codex/gpt-5.3-codex",
    codex_token: str | None = None,
    openai_key: str | None = None,
) -> AsyncGenerator[str, None]:
    """Call LLM and yield raw text response tokens. Supports codex/* and openai/* models."""
    if model_id.startswith("openai/"):
        actual_model = model_id.replace("openai/", "")
        from common.llm_factory import LLMFactory

        if not openai_key:
            raise RuntimeError("OpenAI API key không tồn tại. Thêm Token trong Cài Đặt (Sidebar).")

        factory = LLMFactory()
        model = factory.get_model(model_name=actual_model, streaming=True, api_key=openai_key)
        from langchain_core.messages import HumanMessage, SystemMessage

        messages = [SystemMessage(content=system_prompt), HumanMessage(content=user_content)]
        async for chunk in model.astream(messages):
            if chunk.content:
                yield chunk.content

    else:
        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 = codex_token or get_codex_access_token()
        if not access_token:
            raise RuntimeError("Codex auth not available. Thêm Token trong Cài Đặt (Sidebar).")

        actual_model = model_id.replace("codex/", "") if model_id.startswith("codex/") else model_id

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

        messages = [SystemMessage(content=system_prompt), HumanMessage(content=user_content)]
        async for chunk in model.astream(messages):
            if chunk.content:
                yield chunk.content


# ─── Response Parsing ────────────────────────────────────────────────

def parse_json(raw: str) -> dict:
    """Parse JSON from LLM response, stripping markdown fences. Resilient to incomplete chunks."""
    if not isinstance(raw, str):
        return {"action": "reflect", "data_sufficient": False, "error": "Invalid response"}

    json_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', raw)
    # If no closing ticket is found, try to match open string
    if not json_match:
        open_match = re.search(r'```(?:json)?\s*([\s\S]*)', raw)
        json_str = open_match.group(1).strip() if open_match else raw.strip()
    else:
        json_str = json_match.group(1).strip()

    if not json_str:
        return {"action": "reflect", "data_sufficient": False, "error": "Empty response"}

    try:
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        logger.error("JSON parse error %s on string:\n%s", e, json_str[:200])
        # Return a safe fallback to prevent entire process crash
        return {
            "action": "reflect",
            "data_sufficient": False,
            "error": f"JSON parse error: {str(e)}"
        }

class ThinkingStreamer:
    """Helper to cleanly stream ONLY text inside <thinking>...</thinking> tags."""
    def __init__(self, writer):
        self.writer = writer
        self.buffer = ""
        self.is_thinking = False
        self.done_thinking = False

    def feed(self, token: str):
        if self.done_thinking:
            return

        self.buffer += token

        if not self.is_thinking:
            start_tag = "<thinking>"
            if start_tag in self.buffer:
                self.is_thinking = True
                parts = self.buffer.split(start_tag, 1)
                self.buffer = parts[1]
                # Fall through to process stream
            else:
                # Keep buffer small in case tag spans across chunks
                if len(self.buffer) > 15:
                    self.buffer = self.buffer[-15:]
                return

        if self.is_thinking:
            end_tag = "</thinking>"
            if end_tag in self.buffer:
                self.is_thinking = False
                self.done_thinking = True
                parts = self.buffer.split(end_tag, 1)
                if parts[0]:
                    self.writer({"type": "thinking_token", "token": parts[0]})
                self.buffer = ""
            else:
                # Stream everything except last 15 chars to catch potential end tag splitting
                if len(self.buffer) > 15:
                    to_send = self.buffer[:-15]
                    self.writer({"type": "thinking_token", "token": to_send})
                    self.buffer = self.buffer[-15:]


def extract_html(raw: str) -> str:
    """Extract HTML from LLM response, stripping markdown fences if present."""
    html_match = re.search(r'```(?:html)?\s*([\s\S]*?)```', raw)
    if html_match:
        return html_match.group(1).strip()
    return raw.strip()


# ─── 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:  # NaN
            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) -> dict[str, Any]:
    """Execute a single tool and return results."""
    if tool_name in ("sql_langfuse", "sql_starrocks", "sql_clevertap"):
        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": []}

    elif 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": []}

    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}

    elif tool_name == "read_report_outline":
        report_id = params.get("report_id")
        if not report_id:
            return {"error": "report_id is required"}
        row = await load_report_context(int(report_id))
        if not row or not row.get("html_content"):
            return {"error": f"Report #{report_id} not found or has no content"}
        return {"result": _extract_report_outline(row["html_content"])}

    elif tool_name == "read_report_section":
        report_id = params.get("report_id")
        section = params.get("section", "")
        if not report_id or not section:
            return {"error": "report_id and section are required"}
        row = await load_report_context(int(report_id))
        if not row or not row.get("html_content"):
            return {"error": f"Report #{report_id} not found or has no content"}
        return {"result": _extract_report_section(row["html_content"], section)}

    return {"error": f"Unknown tool: {tool_name}"}


async def execute_tools_parallel(tools: list[dict]) -> list[dict]:
    """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 | 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["preview"] = data[:3] 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.

    Strategy:
    - Small results (≤20 rows): tab-separated table
    - Large results (>20 rows): auto-compute summary statistics
    """
    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)"]

            # Show first 10 rows
            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:")
            date_col = None  # Track date column for trend analysis
            for col in columns:
                values = [row.get(col) for row in data if row.get(col) is not None]
                if not values:
                    continue

                # Detect date columns for trend analysis
                sample = str(values[0]) if values else ""
                is_date = any(kw in col.lower() for kw in ("date", "ngay", "day", "month", "week", "time", "created", "updated"))
                if is_date and not date_col:
                    date_col = col

                try:
                    nums = [float(v) for v in values if v != "" and v is not None]
                    if nums and len(nums) > len(values) * 0.5:
                        sorted_nums = sorted(nums)
                        n = len(sorted_nums)
                        median = sorted_nums[n // 2] if n % 2 else (sorted_nums[n // 2 - 1] + sorted_nums[n // 2]) / 2
                        p25 = sorted_nums[n // 4]
                        p75 = sorted_nums[3 * n // 4]
                        lines.append(
                            f"  {col}: min={min(nums):.2f}, max={max(nums):.2f}, "
                            f"avg={sum(nums)/n:.2f}, median={median:.2f}, "
                            f"P25={p25:.2f}, P75={p75:.2f}, "
                            f"sum={sum(nums):.2f}, count={n}"
                        )
                except (ValueError, TypeError):
                    # Distribution with counts for text columns
                    from collections import Counter
                    counter = Counter(str(v) for v in values)
                    total = len(values)
                    top_items = counter.most_common(5)
                    dist_parts = [f"{val}({cnt}, {cnt*100//total}%)" for val, cnt in top_items]
                    lines.append(
                        f"  {col}: {len(counter)} unique, "
                        f"top: {', '.join(dist_parts)}"
                    )

            # Trend analysis: group numeric values by date column
            if date_col:
                lines.append("### Trend:")
                # Find the first numeric column that's not the date col
                num_cols = []
                for col in columns:
                    if col == date_col:
                        continue
                    try:
                        test_vals = [float(row.get(col, 0)) for row in data[:5] if row.get(col) is not None]
                        if test_vals:
                            num_cols.append(col)
                    except (ValueError, TypeError):
                        pass

                if num_cols:
                    # Group by date, sum numeric columns
                    from collections import defaultdict
                    grouped: dict[str, dict[str, float]] = defaultdict(lambda: {c: 0.0 for c in num_cols})
                    for row in data:
                        date_key = str(row.get(date_col, "unknown"))[:10]  # Trim to date only
                        for nc in num_cols:
                            try:
                                grouped[date_key][nc] += float(row.get(nc, 0) or 0)
                            except (ValueError, TypeError):
                                pass

                    sorted_dates = sorted(grouped.keys())
                    # Show trend for first 2 numeric columns max
                    for nc in num_cols[:2]:
                        vals_by_date = [(d, grouped[d][nc]) for d in sorted_dates]
                        if len(vals_by_date) >= 2:
                            first_val = vals_by_date[0][1]
                            last_val = vals_by_date[-1][1]
                            growth = ((last_val - first_val) / first_val * 100) if first_val != 0 else 0
                            trend_dir = "📈" if growth > 0 else "📉" if growth < 0 else "→"
                            lines.append(
                                f"  {nc} by {date_col}: "
                                f"{vals_by_date[0][0]}={vals_by_date[0][1]:.0f} → "
                                f"{vals_by_date[-1][0]}={vals_by_date[-1][1]:.0f} "
                                f"({trend_dir} {growth:+.1f}%)"
                            )
                else:
                    lines.append(f"  (date column '{date_col}' found but no numeric columns to trend)")

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

    return "\n".join(parts)


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

def humanize_llm_error(exc: Exception) -> str:
    """Convert LLM exceptions to human-readable 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 report generation 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)


# ─── Temp Report Storage ────────────────────────────────────────────

_TEMP_REPORT_ID_SEQ = count(-1, -1)
_TEMP_REPORTS: dict[int, dict[str, Any]] = {}


def create_temp_report(
    *,
    title: str,
    prompt: str,
    html_content: str,
    tools_used: list[str],
    cycles_count: int,
    generation_time_ms: int,
    parent_report_id: int | None = None,
) -> dict[str, Any]:
    """Create a temporary in-memory report when DB save fails."""
    report_id = next(_TEMP_REPORT_ID_SEQ)
    now = datetime.now(timezone.utc).isoformat()
    report = {
        "id": report_id,
        "title": title,
        "prompt": prompt,
        "html_content": html_content,
        "tools_used": tools_used,
        "cycles_count": cycles_count,
        "generation_time_ms": generation_time_ms,
        "report_pages": html_content.count('<div class="page">'),
        "status": "temp",
        "created_by": "temp-memory",
        "created_at": now,
        "updated_at": now,
        "parent_report_id": parent_report_id,
        "persisted": False,
    }
    _TEMP_REPORTS[report_id] = report
    return report


def get_temp_reports() -> dict[int, dict[str, Any]]:
    """Get all temp reports (for listing)."""
    return _TEMP_REPORTS


def delete_temp_report(report_id: int) -> bool:
    """Delete a temp report. Returns True if found and deleted."""
    if report_id in _TEMP_REPORTS:
        del _TEMP_REPORTS[report_id]
        return True
    return False


# ─── Report HTML Parsing Helpers ─────────────────────────────────────

def _strip_html(html: str) -> str:
    """Remove HTML tags and collapse whitespace."""
    text = re.sub(r'<script[^>]*>[\s\S]*?</script>', '', html)
    text = re.sub(r'<style[^>]*>[\s\S]*?</style>', '', text)
    text = re.sub(r'<[^>]+>', ' ', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text


def _extract_report_outline(html: str) -> str:
    """Extract headings, KPIs, and section structure from report HTML.

    Returns a compact outline (~200 words) for the AI to understand
    what the previous report covers without reading full HTML.
    """
    lines = []

    # Extract all headings (h1-h4)
    headings = re.findall(r'<h([1-4])[^>]*>(.*?)</h\1>', html, re.IGNORECASE | re.DOTALL)
    for level, text in headings:
        clean = _strip_html(text).strip()
        if clean:
            indent = "  " * (int(level) - 1)
            lines.append(f"{indent}- H{level}: {clean[:80]}")

    # Extract KPI-like numbers from <strong> tags and table cells
    strong_texts = re.findall(r'<strong[^>]*>(.*?)</strong>', html, re.DOTALL)
    kpis = []
    for s in strong_texts:
        clean = _strip_html(s).strip()
        # Only keep items that contain numbers (likely KPIs)
        if re.search(r'\d', clean) and len(clean) < 100:
            kpis.append(clean)
    if kpis:
        lines.append(f"\n- KPIs found: {', '.join(kpis[:10])}")

    # Count charts (canvas elements = Chart.js)
    chart_count = len(re.findall(r'<canvas', html, re.IGNORECASE))
    if chart_count:
        lines.append(f"- Charts: {chart_count} Chart.js visualizations")

    # Count tables
    table_count = len(re.findall(r'<table', html, re.IGNORECASE))
    if table_count:
        lines.append(f"- Tables: {table_count} data tables")

    # Count pages
    page_count = len(re.findall(r'class="page"', html))
    if page_count:
        lines.append(f"- Pages: {page_count} pages total")

    if not lines:
        return "(empty report — no headings or KPIs found)"

    return "\n".join(lines)


def _extract_report_section(html: str, section_query: str) -> str:
    """Extract the text content of a specific section matching the query.

    Finds the heading that best matches `section_query`, then returns
    all text content until the next heading of the same or higher level.
    """
    # Find all headings with their positions
    heading_pattern = re.compile(r'<h([1-4])[^>]*>(.*?)</h\1>', re.IGNORECASE | re.DOTALL)
    headings = [(m.start(), m.end(), int(m.group(1)), _strip_html(m.group(2)).strip())
                for m in heading_pattern.finditer(html)]

    if not headings:
        return "(no sections found in report)"

    # Find best matching heading (fuzzy match)
    query_lower = section_query.lower()
    best_idx = -1
    best_score = 0
    for i, (start, end, level, text) in enumerate(headings):
        text_lower = text.lower()
        # Exact substring match gets highest score
        if query_lower in text_lower or text_lower in query_lower:
            score = 100
        else:
            # Word overlap score
            query_words = set(query_lower.split())
            text_words = set(text_lower.split())
            common = query_words & text_words
            score = len(common) * 10 if common else 0

        if score > best_score:
            best_score = score
            best_idx = i

    if best_idx == -1:
        available = [h[3] for h in headings[:10]]
        return f"Section '{section_query}' not found. Available: {', '.join(available)}"

    # Extract content from this heading to the next heading of same or higher level
    match_start, match_end, match_level, match_title = headings[best_idx]
    content_start = match_end

    # Find the end: next heading of same or higher level
    content_end = len(html)
    for i in range(best_idx + 1, len(headings)):
        next_start, _, next_level, _ = headings[i]
        if next_level <= match_level:
            content_end = next_start
            break

    section_html = html[content_start:content_end]
    section_text = _strip_html(section_html)

    # Truncate to ~1500 chars to avoid token explosion
    if len(section_text) > 1500:
        section_text = section_text[:1500] + "... (truncated)"

    return f"## {match_title}\n{section_text}"


async def load_report_context(report_id: int) -> dict[str, Any] | None:
    """Load report context from temp memory or DB."""
    if report_id in _TEMP_REPORTS:
        return _TEMP_REPORTS[report_id]

    import psycopg
    from config import CHECKPOINT_POSTGRES_URL
    from common.pool_wrapper import get_pooled_connection_compat

    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute(
            "SELECT id, title, prompt, tools_used, html_content, created_at "
            "FROM dashboard_canifa.saved_reports WHERE id = %s",
            (report_id,),
        )
        row = cur.fetchone()
        cur.close()
        if row:
            return {
                "id": row[0], "title": row[1], "prompt": row[2],
                "tools_used": row[3], "html_content": row[4], "created_at": row[5],
            }
        return None
    except Exception as e:
        logger.warning("Failed to load report context from DB for %s: %s", report_id, e)
        return None
    finally:
        if conn: conn.close()


# ─── SSE Helper ──────────────────────────────────────────────────────

def sse_event(data: dict) -> str:
    """Format dict as SSE event string."""
    return f"data: {json.dumps(data, ensure_ascii=False, default=str)}\n\n"
