"""
AI Data Analyst — Text-to-SQL Chat API.

POST /api/sql-chat
  Input: { question: str, history: list[{role, content}] }
  Output: { sql: str, data: list, columns: list, chart_type: str, x_key, y_keys, explanation: str }

Uses Codex (ChatGPT backend API) to generate SQL from Vietnamese questions,
executes on StarRocks, and returns structured results with optional Chart.js config.
"""

import json
import logging
import re
import traceback
from typing import Any

from fastapi import APIRouter, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel

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

logger = logging.getLogger(__name__)
router = APIRouter()

TABLE_NAME = "shared_source.magento_product_dimension_with_text_embedding"

# ─── Database Schema (provided to LLM as context) ────────────────────

DB_SCHEMA = f"""
## Database: StarRocks
## Table: {TABLE_NAME}

### Columns (SELECT only, KHÔNG query cột 'vector'):
- internal_ref_code (VARCHAR) — Mã sản phẩm nội bộ, VD: "8TP25A005"
- magento_ref_code (VARCHAR) — Mã ref = internal_ref_code + color, VD: "8TP25A005-SW011"
- product_color_code (VARCHAR) — Mã màu sản phẩm
- product_name (VARCHAR) — Tên sản phẩm, VD: "ÁO THUN NAM"
- color_code (VARCHAR) — Mã màu
- master_color (VARCHAR) — Nhóm màu chính: "Trắng", "Đen", "Đỏ", etc.
- product_color_name (VARCHAR) — Tên màu đầy đủ
- season_sale (VARCHAR) — Mùa sale
- season (VARCHAR) — Mùa: "Xuân Hè", "Thu Đông", etc.
- style (VARCHAR) — Phong cách sản phẩm
- fitting (VARCHAR) — Kiểu form: "Regular Fit", "Slim Fit", etc.
- size_scale (VARCHAR) — Thang size: "XS-XL", "1-6", etc.
- graphic (VARCHAR) — Họa tiết
- pattern (VARCHAR) — Hoa văn
- weaving (VARCHAR) — Kiểu dệt
- shape_detail (VARCHAR) — Chi tiết kiểu dáng
- form_neckline (VARCHAR) — Kiểu cổ: "Cổ tròn", "Cổ V", etc.
- form_sleeve (VARCHAR) — Kiểu tay: "Ngắn tay", "Dài tay", etc.
- form_length (VARCHAR) — Kiểu dài
- form_waistline (VARCHAR) — Kiểu eo
- form_shoulderline (VARCHAR) — Kiểu vai
- material (VARCHAR) — Chất liệu
- product_group (VARCHAR) — Nhóm sản phẩm
- product_line_vn (VARCHAR) — Dòng sản phẩm VN: "Áo Thun", "Quần Jean", "Váy/Đầm", etc.
- unit_of_measure (VARCHAR) — Đơn vị
- sale_price (DECIMAL) — Giá bán
- original_price (DECIMAL) — Giá gốc
- discount_amount (DECIMAL) — Số tiền giảm giá (= original_price - sale_price)
- material_group (VARCHAR) — Nhóm chất liệu
- product_line_en (VARCHAR) — Dòng sản phẩm EN
- age_by_product (VARCHAR) — Nhóm tuổi
- gender_by_product (VARCHAR) — Giới tính: "nam", "nữ", "unisex", "bé trai", "bé gái"
- quantity_sold (DECIMAL) — Số lượng đã bán
- is_new_product (TINYINT) — Sản phẩm mới: 1 = mới, 0 = cũ
- product_image_url (VARCHAR) — URL ảnh sản phẩm
- description_text (VARCHAR) — Mô tả ngắn
- product_image_url_thumbnail (VARCHAR) — URL ảnh thumbnail
- product_web_url (VARCHAR) — URL trang sản phẩm
- product_web_material (VARCHAR) — Chất liệu trên web
- description_text_full (VARCHAR) — Mô tả đầy đủ
## NOTE: KHÔNG query cột 'vector' (quá lớn, dùng cho embedding)
"""

SYSTEM_PROMPT = f"""Bạn là AI Data Analyst cho Canifa (thương hiệu thời trang Việt Nam).
User sẽ hỏi bằng tiếng Việt về dữ liệu sản phẩm. Bạn cần:

1. Phân tích câu hỏi
2. Viết SQL query phù hợp
3. Chọn loại biểu đồ phù hợp nhất từ 25 loại

{DB_SCHEMA}

## 25 CHART TYPES:
| Type | Dùng khi |
|---|---|
| bar | So sánh categories (top sản phẩm, so sánh dòng) |
| horizontalBar | Bars ngang (tên dài, ranking) |
| multiBar | So sánh 2+ metrics cùng lúc |
| stackedBar | Tỷ lệ composition theo category |
| line | Trend theo thời gian |
| area | Volume theo thời gian (fill dưới line) |
| stackedArea | Composition thay đổi theo thời gian |
| pie | Phân bổ % (< 7 categories) |
| doughnut | Phân bổ % với tổng ở giữa |
| scatter | Correlation 2 biến (giá vs lượng bán) |
| combo | Bar + Line cùng lúc (VD: doanh thu + growth%) |
| bubble | 3 biến (x, y, size) |
| radar | So sánh đa chiều (spider chart) |
| polarArea | Phân bổ weighted |
| waterfall | Sequential tăng/giảm (breakdown) |
| numberCard | Chỉ 1 con số KPI lớn (tổng SP, avg giá) |
| gauge | Tiến độ % (% đạt target) |
| progress | Multiple progress bars (so sánh tiến độ) |
| funnel | Conversion pipeline (phễu) |
| table | Dữ liệu dạng bảng (không cần chart) |
| treemap | Hierarchical blocks |
| heatmap | Color-coded matrix |
| boxplot | Distribution statistics |
| sankey | Flow diagram |
| pivotTable | Cross-tab grouped table |

## RULES:
- CHỈNH sử dụng SELECT. TUYỆT ĐỐI KHÔNG dùng INSERT, UPDATE, DELETE, DROP, ALTER, CREATE.
- Luôn thêm LIMIT (tối đa 200 rows)
- KHÔNG query cột 'embedding' (quá lớn)
- Trả lời JSON format theo schema bên dưới

## OUTPUT FORMAT (trả về JSON):
```json
{{
  "thinking": "Phân tích ngắn gọn câu hỏi",
  "sql": "SELECT ... FROM {TABLE_NAME} WHERE ... LIMIT 50",
  "chart_type": "bar|line|pie|doughnut|numberCard|...",
  "x_key": "tên cột làm trục X hoặc labels",
  "y_keys": ["tên cột(s) làm giá trị Y"],
  "y_labels": {{"column_name": "Label hiển thị"}},
  "trend": {{"value": 12.5, "direction": "up", "label": "vs tháng trước"}} hoặc null,
  "total_label": "Tổng sản phẩm" (cho doughnut),
  "gauge_max": 100 (cho gauge),
  "footer": "Ghi chú thêm",
  "explanation": "Giải thích kết quả bằng tiếng Việt"
}}
```

## TIPS CHỌN chart_type:
- Câu hỏi "có bao nhiêu" / "tổng" → numberCard
- Câu hỏi "top N" → bar hoặc horizontalBar
- Câu hỏi "phân bổ" / "tỷ lệ" → pie hoặc doughnut
- Câu hỏi "so sánh" 2+ metrics → multiBar hoặc combo
- Câu hỏi "trend" / "theo thời gian" → line hoặc area
- Câu hỏi "ranking" → horizontalBar hoặc progress
- Câu hỏi conversion → funnel
- Câu hỏi matrix / cross-tab → heatmap hoặc pivotTable
"""


# ─── Request/Response Models ─────────────────────────────────────────

class ChatMessage(BaseModel):
    role: str  # "user" or "assistant"
    content: str

class SqlChatRequest(BaseModel):
    question: str
    history: list[ChatMessage] = []
    model: str = "codex/gpt-5.3-codex"

class SqlChatResponse(BaseModel):
    sql: str = ""
    data: list[dict] = []
    columns: list[str] = []
    chart_type: str | None = None

    explanation: str = ""
    row_count: int = 0
    error: str | None = None


# ─── 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) -> 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()}"
    # Block selecting the vector/embedding column (too large)
    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"
    return None


# ─── Codex LLM Call ──────────────────────────────────────────────────

async def _call_codex(question: str, history: list[ChatMessage], model_id: str = "codex/gpt-5.3-codex") -> dict:
    """Call Codex LLM to generate SQL from user question."""
    from common.codex_auth import get_codex_access_token, get_codex_account_id, get_refresh_token

    access_token = get_codex_access_token()
    if not access_token:
        raise HTTPException(status_code=503, detail="Codex auth not available. Run `codex login` first.")

    account_id = get_codex_account_id() or ""
    refresh_token = get_refresh_token()

    # Strip codex/ prefix to get actual model name for the API
    actual_model = model_id.replace("codex/", "") if model_id.startswith("codex/") else model_id
    logger.info("🔧 Using Codex model: %s (from %s)", actual_model, model_id)

    from common.codex_chat_model import CodexChatModel
    from langchain_core.messages import HumanMessage, SystemMessage, AIMessage

    model = CodexChatModel(
        model_name=actual_model,
        access_token=access_token,
        account_id=account_id,
        refresh_token=refresh_token,
        streaming=False,
    )

    # Build messages
    messages = [SystemMessage(content=SYSTEM_PROMPT)]
    for msg in history[-6:]:  # Keep last 6 messages for context
        if msg.role == "user":
            messages.append(HumanMessage(content=msg.content))
        else:
            messages.append(AIMessage(content=msg.content))
    messages.append(HumanMessage(content=question))

    # Call LLM
    result = await model._agenerate(messages)
    response_text = result.generations[0].message.content

    # Parse JSON from response
    # Try to extract JSON from markdown code block or raw text
    json_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', response_text)
    if json_match:
        json_str = json_match.group(1).strip()
    else:
        json_str = response_text.strip()

    try:
        parsed = json.loads(json_str)
    except json.JSONDecodeError:
        logger.warning("Failed to parse LLM JSON response: %s", response_text[:500])
        # Fallback: try to extract SQL from response
        sql_match = re.search(r'(SELECT\s+.+?)(?:;|$)', response_text, re.IGNORECASE | re.DOTALL)
        parsed = {
            "sql": sql_match.group(1).strip() if sql_match else "",
            "chart_type": None,
            "x_key": None, "y_keys": None, "y_labels": None,
            "trend": None, "total_label": None, "gauge_max": None, "footer": None,
            "explanation": response_text[:200],
        }

    return parsed


# ─── API Endpoint ────────────────────────────────────────────────────

@router.post("/api/sql-chat", summary="AI Text-to-SQL Chat")
async def sql_chat(req: SqlChatRequest):
    """Convert Vietnamese question to SQL, execute, return results + chart."""
    try:
        # 1. Call Codex LLM
        logger.info("🤖 SQL Chat: %s", req.question[:100])
        llm_result = await _call_codex(req.question, req.history, req.model)

        sql = llm_result.get("sql", "").strip()
        explanation = llm_result.get("explanation", "")
        thinking = llm_result.get("thinking", "")

        chart_type = llm_result.get("chart_type")
        x_key = llm_result.get("x_key")
        y_keys = llm_result.get("y_keys")
        y_labels = llm_result.get("y_labels")
        trend = llm_result.get("trend")
        total_label = llm_result.get("total_label")
        gauge_max = llm_result.get("gauge_max")
        footer = llm_result.get("footer")

        # Build chart metadata dict
        chart_meta = {
            "chart_type": chart_type,
            "x_key": x_key,
            "y_keys": y_keys,
            "y_labels": y_labels,
            "trend": trend,
            "total_label": total_label,
            "gauge_max": gauge_max,
            "footer": footer,
        }

        if not sql:
            return JSONResponse(content={
                "sql": "",
                "data": [],
                "columns": [],
                **chart_meta,
                "explanation": explanation or "Không thể tạo SQL query cho câu hỏi này.",
                "row_count": 0,
                "error": None,
            })

        # 2. Validate SQL safety
        safety_error = _validate_sql(sql)
        if safety_error:
            return JSONResponse(content={
                "sql": sql,
                "data": [],
                "columns": [],
                "chart_type": None,
                "explanation": f"⚠️ SQL không hợp lệ: {safety_error}",
                "row_count": 0,
                "error": safety_error,
            })

        # 3. Execute SQL on StarRocks
        db = StarRocksConnection()
        try:
            rows = await db.execute_query_async(sql)
        except Exception as e:
            error_msg = str(e)
            logger.error("SQL execution error: %s\nSQL: %s", error_msg, sql)
            return JSONResponse(content={
                "sql": sql,
                "data": [],
                "columns": [],
                "chart_type": None,
                "explanation": f"❌ Lỗi thực thi SQL: {error_msg[:200]}",
                "row_count": 0,
                "error": error_msg[:200],
            })

        # 4. Format response
        columns = list(rows[0].keys()) if rows else []
        data = [dict(row) for row in rows]

        # Convert non-serializable types
        from decimal import Decimal
        for row in data:
            for key, val in row.items():
                if isinstance(val, Decimal):
                    row[key] = float(val)
                elif isinstance(val, (bytes, bytearray)):
                    row[key] = val.hex()
                elif hasattr(val, 'isoformat'):
                    row[key] = val.isoformat()

        logger.info("✅ SQL Chat result: %d rows, %d columns", len(data), len(columns))

        return JSONResponse(content={
            "sql": sql,
            "data": data,
            "columns": columns,
            **chart_meta,
            "explanation": explanation,
            "row_count": len(data),
            "error": None,
        })

    except HTTPException:
        raise
    except Exception as e:
        logger.error("SQL Chat error: %s\n%s", e, traceback.format_exc())
        return JSONResponse(
            status_code=500,
            content={"error": str(e)[:200], "sql": "", "data": [], "columns": [], "explanation": ""},
        )


@router.get("/api/sql-chat/status", summary="Check Codex availability")
async def sql_chat_status():
    """Check if Codex auth is available."""
    from common.codex_auth import is_codex_available
    return {"codex_available": is_codex_available()}


# ═══════════════════════════════════════════════════════════════════════
# DashboardAI — Multi-Widget Dashboard Generator (SSE Streaming)
# ═══════════════════════════════════════════════════════════════════════

import asyncio
from starlette.responses import StreamingResponse
from agent.sql_agent.prompts.dashboard_prompt import DASHBOARD_PROMPT, STARROCKS_TABLE as DASH_TABLE


class DashboardRequest(BaseModel):
    question: str
    history: list[ChatMessage] = []
    model: str = "codex/gpt-5.3-codex"


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


async def _call_dashboard_llm(question: str, history: list, model_name: str) -> dict:
    """Call Codex LLM to generate dashboard JSON."""
    from common.codex_auth import get_codex_access_token, get_codex_account_id, get_refresh_token
    access_token = get_codex_access_token()
    if not access_token:
        raise HTTPException(status_code=503, detail="Codex auth not available.")

    account_id = get_codex_account_id() or ""
    refresh_token = get_refresh_token()
    actual_model = model_name.replace("codex/", "") if model_name.startswith("codex/") else model_name
    logger.info("🔧 Dashboard model: %s", actual_model)

    from common.codex_chat_model import CodexChatModel
    from langchain_core.messages import HumanMessage, SystemMessage, AIMessage

    model = CodexChatModel(
        model_name=actual_model, access_token=access_token,
        account_id=account_id, refresh_token=refresh_token, streaming=False,
    )

    messages = [SystemMessage(content=DASHBOARD_PROMPT)]
    for msg in history[-4:]:
        if msg.role == "user":
            messages.append(HumanMessage(content=msg.content))
        else:
            messages.append(AIMessage(content=msg.content))
    messages.append(HumanMessage(content=question))

    result = await model._agenerate(messages)
    response_text = result.generations[0].message.content

    json_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', response_text)
    json_str = json_match.group(1).strip() if json_match else response_text.strip()
    return json.loads(json_str)


async def _exec_widget_sql(db_starrocks, db_postgres, widget: dict) -> dict:
    """Execute a single widget's SQL query and attach results.
    Routes to StarRocks or PostgreSQL based on widget['db'] field.
    """
    sql = widget.get("sql", "")
    if not sql:
        widget["data"] = []
        widget["columns"] = []
        widget["error"] = None
        return widget

    safety_err = _validate_sql(sql)
    if safety_err:
        widget["data"] = []
        widget["columns"] = []
        widget["error"] = safety_err
        return widget

    # Route to correct database
    target_db = widget.get("db", "starrocks")

    try:
        if target_db == "postgres":
            rows = await db_postgres.execute_query_async(sql)
            columns = list(rows[0].keys()) if rows else []
            data = [_serialize_row(dict(r)) for r in rows]
        else:
            rows = await db_starrocks.execute_query_async(sql)
            columns = list(rows[0].keys()) if rows else []
            data = [_serialize_row(dict(r)) for r in rows]

        widget["data"] = data
        widget["columns"] = columns
        widget["error"] = None
        logger.info("  ✅ Widget %s [%s]: %d rows", widget.get("id", "?"), target_db, len(data))
    except Exception as e:
        logger.error("  ❌ Widget %s [%s] error: %s", widget.get("id", "?"), target_db, str(e)[:200])
        widget["data"] = []
        widget["columns"] = []
        widget["error"] = str(e)[:200]

    return widget


@router.post("/api/sql-dashboard", summary="AI Dashboard Generator (SSE Stream)")
async def sql_dashboard_stream(req: DashboardRequest):
    """Generate a dashboard via SSE — streams header first, then each widget as SQL completes."""

    async def event_stream():
        try:
            logger.info("📊 Dashboard AI (SSE): %s", req.question[:100])

            # Phase 1: Call LLM → get dashboard layout
            dashboard = await _call_dashboard_llm(req.question, req.history, req.model)
            widgets = dashboard.get("widgets", [])
            logger.info("📊 Dashboard: %s — %d widgets", dashboard.get("title", "?"), len(widgets))

            # Send header (title + subtitle + widget skeletons without data)
            header = {
                "type": "header",
                "title": dashboard.get("title", "Dashboard"),
                "subtitle": dashboard.get("subtitle", ""),
                "widgets": [
                    {"id": w.get("id"), "type": w.get("type"), "title": w.get("title"),
                     "size": w.get("size", "md"), "color": w.get("color", "indigo")}
                    for w in widgets
                ],
            }
            yield f"data: {json.dumps(header, ensure_ascii=False)}\n\n"

            # Phase 2: Execute all SQLs concurrently, stream each as it completes
            db_starrocks = StarRocksConnection()
            db_postgres = PostgresReadonly()
            tasks = {
                asyncio.create_task(_exec_widget_sql(db_starrocks, db_postgres, w)): w.get("id")
                for w in widgets
            }

            for done_task in asyncio.as_completed(tasks):
                widget = await done_task
                event = {
                    "type": "widget",
                    "id": widget.get("id"),
                    "data": widget.get("data", []),
                    "columns": widget.get("columns", []),
                    "x_key": widget.get("x_key"),
                    "y_key": widget.get("y_key"),
                    "y_keys": widget.get("y_keys"),
                    "error": widget.get("error"),
                }
                yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"

            # Done signal
            yield f"data: {json.dumps({'type': 'done'})}\n\n"

        except json.JSONDecodeError:
            yield f"data: {json.dumps({'type': 'error', 'message': 'AI response was not valid JSON'})}\n\n"
        except Exception as e:
            logger.error("Dashboard SSE error: %s\n%s", e, traceback.format_exc())
            yield f"data: {json.dumps({'type': 'error', 'message': str(e)[:200]})}\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")


