"""
AI Report HTML Generator — SSE Streaming Endpoint.

POST /api/report-html
  Input:  { question: str }
  Output: SSE stream → final event contains body HTML

Thin API layer that delegates to:
  - agent/report_agent/main_graph.py  (report generation)
  - agent/report_agent/inline_graph.py (inline editing)
"""

import json
import logging
import time
import traceback
import uuid
from typing import Any

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

from fastapi import APIRouter, Depends
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from starlette.responses import StreamingResponse

from auth.middleware import get_current_user
from agent.report_agent.core import (
    humanize_llm_error,
    create_temp_report,
    get_temp_reports,
    delete_temp_report,
    sse_event as _sse,
)
from agent.report_agent.main_graph import run_report_agent
from agent.report_agent.inline_graph import run_inline_agent
from agent.report_agent.follow_up_graph import run_followup_agent


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


# ─── Request Models ──────────────────────────────────────────────────

class ReportHTMLRequest(BaseModel):
    question: str
    model: str = "codex/gpt-5.3-codex"
    parent_report_id: int | None = None
    current_html: str | None = None
    conversation_id: str | None = None


# ─── SSE Streaming Endpoint ─────────────────────────────────────────

@router.post("/api/report-html", summary="AI Report HTML Generator (SSE Stream)")
async def report_html_stream(req: ReportHTMLRequest, user: dict = Depends(get_current_user)):
    """
    Generate a report as raw HTML body content via SSE.
    Strategy:
      1. Try to enqueue via Redis → return task_id (client polls /stream/{task_id})
      2. Fallback to direct SSE if Redis unavailable
    """
    user_codex_token = user.get("settings", {}).get("codex_token")
    user_openai_key = user.get("settings", {}).get("openai_key")
    conv_id = req.conversation_id or str(uuid.uuid4())

    # ── Strategy 1: Redis Queue (async, resumable) ──────────────────
    from agent.report_agent.report_queue import enqueue_report_task
    task_id = await enqueue_report_task(
        question=req.question,
        model=req.model,
        conversation_id=conv_id,
        parent_report_id=req.parent_report_id,
        current_html=req.current_html,
        codex_token=user_codex_token,
        openai_key=user_openai_key,
    )
    if task_id:
        # Return task_id — client will connect to /stream/{task_id}
        return JSONResponse({
            "mode": "queued",
            "task_id": task_id,
            "conversation_id": conv_id,
            "stream_url": f"/api/report-html/stream/{task_id}",
        })

    # ── Strategy 2: Direct SSE fallback (old behavior) ──────────────
    logger.warning("Redis unavailable, falling back to direct SSE")

    async def event_stream():
        # Save user message to chat_history
        conn_ct = None
        try:
            conn_ct = get_pooled_connection_compat()
            cur_ct = conn_ct.cursor()
            cur_ct.execute(
                """INSERT INTO dashboard_canifa.chat_history
                   (user_id, module, conversation_id, message, is_human)
                   VALUES (1, 'report_generator', %s, %s, true)""",
                (conv_id, req.question)
            )
            cur_ct.close()
        except Exception as e:
            logger.error("Failed to save user msg to chat_history: %s", e)
        finally:
            if conn_ct: conn_ct.close()

        start_time = time.time()
        try:
            logger.info("📝 HTML Report Agent (direct SSE): %s", req.question[:100])

            html_body = ""
            tools_used: list[str] = []
            cycles_count = 0

            async for event in run_report_agent(
                question=req.question,
                model=req.model,
                codex_token=user_codex_token,
                openai_key=user_openai_key,
                parent_report_id=req.parent_report_id,
                current_html=req.current_html,
            ):
                event_type = event.get("type")

                if event_type == "direct_response":
                    message = event.get("message", "")
                    yield _sse(event)
                    _save_chat_response(conv_id, message)
                    continue

                if event_type == "html_complete":
                    html_body = event["html"]
                    tools_used = event.get("tools_used", [])
                    cycles_count = event.get("cycles_count", 0)

                    generation_time_ms = int((time.time() - start_time) * 1000)
                    report_id, persisted, final_html = await _save_report(
                        question=req.question,
                        html_body=html_body,
                        tools_used=tools_used,
                        cycles_count=cycles_count,
                        generation_time_ms=generation_time_ms,
                        parent_report_id=req.parent_report_id,
                        current_html=req.current_html,
                        conv_id=conv_id,
                    )

                    yield _sse({
                        "type": "html_complete",
                        "html": html_body,
                        "report_id": report_id,
                        "conversation_id": conv_id,
                        "persisted": persisted,
                    })
                    continue

                yield _sse(event)

        except json.JSONDecodeError as e:
            logger.error("Report HTML JSON error: %s", e)
            yield _sse({"type": "error", "message": f"AI response was not valid JSON: {str(e)[:200]}"})
        except RuntimeError as e:
            logger.error("Report HTML runtime error: %s", e)
            yield _sse({"type": "error", "message": str(e)[:300]})
        except Exception as e:
            friendly_error = humanize_llm_error(e)
            logger.error("Report HTML error: %s\n%s", e, traceback.format_exc())
            yield _sse({"type": "error", "message": friendly_error[:500]})

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


# ─── Redis Queue SSE Stream Reader ───────────────────────────────────

@router.get("/api/report-html/stream/{task_id}", summary="Stream report events from Redis")
async def stream_report_events(task_id: str, last_id: str = "0-0"):
    """
    SSE endpoint that reads events from Redis Stream.
    Supports reconnection: pass ?last_id=<msg_id> to resume from last position.
    """
    from agent.report_agent.report_queue import read_events, get_task_status

    async def event_stream():
        cursor = last_id
        max_wait_cycles = 600  # 600 × 2s = 20 minutes max wait
        wait_cycle = 0

        while wait_cycle < max_wait_cycles:
            events = await read_events(task_id, last_id=cursor, block_ms=2000)

            if not events:
                # Check if task is done/error (no more events coming)
                status = await get_task_status(task_id)
                if status and status.get("status") in ("done", "error"):
                    break
                wait_cycle += 1
                continue

            wait_cycle = 0  # Reset on activity
            for msg_id, event_data in events:
                cursor = msg_id
                yield _sse({**event_data, "_msg_id": msg_id})

                if event_data.get("type") in ("done", "error"):
                    return

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


# ─── DB Save Helpers ─────────────────────────────────────────────────

def _save_chat_response(conv_id: str, message: str):
    """Save bot response to chat_history."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute(
            """INSERT INTO dashboard_canifa.chat_history
               (user_id, module, conversation_id, message, is_human)
               VALUES (1, 'report_generator', %s, %s, false)""",
            (conv_id, message)
        )
        cur.close()
    except Exception as e:
        logger.error("Failed to save chat response: %s", e)
    finally:
        if conn: conn.close()


async def _save_report(
    *,
    question: str,
    html_body: str,
    tools_used: list[str],
    cycles_count: int,
    generation_time_ms: int,
    parent_report_id: int | None,
    current_html: str | None,
    conv_id: str,
) -> tuple[int, bool, str]:
    """
    Save report to DB. Returns (report_id, persisted, final_html_body).
    Falls back to temp in-memory storage if DB fails.
    """
    version = 1
    parent_id = parent_report_id
    final_html_body = html_body
    title = question[:100] + ("..." if len(question) > 100 else "")
    is_followup = bool(parent_id)

    # Handle follow-up: merge HTML
    if parent_id:
        conn_db = None
        try:
            conn_db = get_pooled_connection_compat()
            cur_db = conn_db.cursor()

            cur_db.execute(
                "SELECT COALESCE(MAX(version), 0) + 1 FROM dashboard_canifa.saved_reports "
                "WHERE parent_id = %s OR id = %s",
                (parent_id, parent_id),
            )
            version = cur_db.fetchone()[0]

            cur_db.execute(
                "SELECT html_content, title FROM dashboard_canifa.saved_reports WHERE id = %s",
                (parent_id,),
            )
            row = cur_db.fetchone()
            if row:
                old_html, old_title = row
                title = old_title
                base_html = current_html if current_html else old_html
                final_html_body = base_html + "\n<hr>\n" + html_body

            cur_db.close()
        except Exception as e:
            logger.error("Error getting parent report data: %s", e)
            version = 2
        finally:
            if conn_db: conn_db.close()

    # Insert report
    conn_db = None
    try:
        conn_db = get_pooled_connection_compat()
        cur_db = conn_db.cursor()
        cur_db.execute(
            """INSERT INTO dashboard_canifa.saved_reports
            (title, prompt, html_content, tools_used, cycles_count, generation_time_ms, version, parent_id, conversation_id)
            VALUES (%s, %s, %s, %s::jsonb, %s, %s, %s, %s, %s)
            RETURNING id""",
            (title, question, final_html_body, json.dumps(tools_used), cycles_count,
             generation_time_ms, version, parent_id, conv_id),
        )
        report_id = cur_db.fetchone()[0]

        # Save AI message referencing the report
        action_label = "tạo" if version == 1 else "cập nhật"
        cur_db.execute(
            """INSERT INTO dashboard_canifa.chat_history
               (user_id, module, conversation_id, message, is_human)
               VALUES (1, 'report_generator', %s, %s, false)""",
            (conv_id, f"Đã {action_label} báo cáo: {title} (v{version})|REPORT:{report_id}")
        )

        cur_db.close()
        logger.info("💾 Saved report to DB ID: %s (v%s)", report_id, version)
        return report_id, True, final_html_body

    except Exception as e:
        logger.error("Failed to save report to DB: %s", e)
        temp_report = create_temp_report(
            title=title,
            prompt=question,
            html_content=final_html_body,
            tools_used=tools_used,
            cycles_count=cycles_count,
            generation_time_ms=generation_time_ms,
            parent_report_id=parent_report_id,
        )
        logger.info("🧠 Saved report to in-memory fallback with ID: %s", temp_report["id"])
        return temp_report["id"], False, final_html_body
    finally:
        if conn_db: conn_db.close()


# ─── Report History API ──────────────────────────────────────────────

@router.get("/reports")
async def get_reports():
    """List all saved reports (metadata only, no full HTML)."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute("""
            SELECT id, title, prompt, tools_used, cycles_count, generation_time_ms, version, parent_id, created_at, conversation_id
            FROM dashboard_canifa.saved_reports
            ORDER BY created_at DESC
            LIMIT 50
        """)
        rows = []
        for r in cur.fetchall():
            rows.append({
                "id": r[0], "title": r[1], "prompt": r[2], "tools_used": r[3],
                "cycles_count": r[4], "generation_time_ms": r[5],
                "version": r[6], "parent_id": r[7], "conversation_id": r[9],
                "created_at": r[8].isoformat() if r[8] else None,
                "persisted": True,
            })
        cur.close()

        # Add in-memory temp reports
        temp_rows = []
        for report in get_temp_reports().values():
            temp_rows.append({
                "id": report["id"], "title": report["title"], "prompt": report["prompt"],
                "tools_used": report["tools_used"], "cycles_count": report["cycles_count"],
                "generation_time_ms": report["generation_time_ms"],
                "version": 1, "parent_id": None,
                "created_at": report["created_at"], "persisted": False,
            })
        combined = sorted([*temp_rows, *rows], key=lambda x: x.get("created_at") or "", reverse=True)[:50]
        return {"reports": combined}
    except Exception as e:
        logger.error("Error fetching reports: %s", e)
        return {"error": str(e), "reports": []}
    finally:
        if conn: conn.close()


@router.get("/api/reports/conversations")
async def get_conversations(user: dict = Depends(get_current_user)):
    """Get list of conversations grouped by conversation_id."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute("""
            WITH LastMsgs AS (
                SELECT conversation_id, message AS title, created_at,
                       ROW_NUMBER() OVER(PARTITION BY conversation_id ORDER BY created_at ASC) as rn,
                       MAX(created_at) OVER(PARTITION BY conversation_id) as last_updated
                FROM dashboard_canifa.chat_history
                WHERE module = 'report_generator' AND is_human = true
            )
            SELECT conversation_id, title, last_updated
            FROM LastMsgs
            WHERE rn = 1
            ORDER BY last_updated DESC LIMIT 50;
        """)
        conversations = []
        for r in cur.fetchall():
            conversations.append({
                "conversation_id": r[0],
                "title": r[1][:80] + "..." if len(r[1]) > 80 else r[1],
                "updated_at": r[2].isoformat() if r[2] else None
            })
        cur.close()
        return {"conversations": conversations}
    except Exception as e:
        logger.error("Error fetching conversations: %s", e)
        return {"error": str(e), "conversations": []}
    finally:
        if conn: conn.close()


@router.get("/api/reports/conversations/{conv_id}")
async def get_conversation_data(conv_id: str, user: dict = Depends(get_current_user)):
    """Get chat history and reports for a specific conversation."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()

        cur.execute("""
            SELECT id, message, is_human, created_at
            FROM dashboard_canifa.chat_history
            WHERE module = 'report_generator' AND conversation_id = %s
            ORDER BY created_at ASC
        """, (conv_id,))
        messages = []
        for r in cur.fetchall():
            messages.append({
                "id": r[0], "text": r[1], "is_human": r[2],
                "created_at": r[3].isoformat() if r[3] else None
            })

        cur.execute("""
            SELECT id, title, prompt, version, parent_id, created_at
            FROM dashboard_canifa.saved_reports
            WHERE conversation_id = %s
            ORDER BY created_at ASC
        """, (conv_id,))
        reports = []
        for r in cur.fetchall():
            reports.append({
                "id": r[0], "title": r[1], "prompt": r[2],
                "version": r[3], "parent_id": r[4],
                "created_at": r[5].isoformat() if r[5] else None
            })

        cur.close()
        return {"conversation_id": conv_id, "messages": messages, "reports": reports}
    except Exception as e:
        logger.error(f"Error fetching conversation {conv_id}: %s", e)
        return {"error": str(e)}
    finally:
        if conn: conn.close()


@router.get("/api/reports/{report_id}")
async def get_report_by_id(report_id: int, user: dict = Depends(get_current_user)):
    """Get a single report with full HTML content."""
    temp_reports = get_temp_reports()
    if report_id in temp_reports:
        return {"report": temp_reports[report_id]}

    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute(
            "SELECT id, title, prompt, html_content, tools_used, version, parent_id, "
            "cycles_count, generation_time_ms, created_at "
            "FROM dashboard_canifa.saved_reports WHERE id = %s",
            (report_id,),
        )
        row = cur.fetchone()
        cur.close()

        if not row:
            return {"error": "Report not found"}

        return {"report": {
            "id": row[0], "title": row[1], "prompt": row[2], "html_content": row[3],
            "tools_used": row[4], "version": row[5], "parent_id": row[6],
            "cycles_count": row[7], "generation_time_ms": row[8],
            "created_at": row[9].isoformat() if row[9] else None,
        }}
    except Exception as e:
        logger.error("Error fetching report %d: %s", report_id, e)
        return {"error": str(e)}
    finally:
        if conn: conn.close()


@router.delete("/api/reports/{report_id}")
async def delete_report(report_id: int, user: dict = Depends(get_current_user)):
    """Delete a generated report."""
    if delete_temp_report(report_id):
        return {"success": True, "status": "DELETED 1 (temp-memory)"}

    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()
        cur.execute("DELETE FROM dashboard_canifa.saved_reports WHERE id = %s", (report_id,))
        cur.close()
        return {"success": True, "status": "DELETED"}
    except Exception as e:
        logger.error("Error deleting report %d: %s", report_id, e)
        return {"error": str(e)}
    finally:
        if conn: conn.close()


# ─── Inline Edit API ────────────────────────────────────────────────

class InlineEditRequest(BaseModel):
    selected_text: str
    action: str = "rewrite"  # rewrite | enrich | shorten | fix | agent_rewrite
    context: str = ""
    model: str = "codex/gpt-5.3-codex"


@router.post("/api/report-inline-edit")
async def inline_edit(req: InlineEditRequest, user: dict = Depends(get_current_user)):
    """AI-powered inline edit — delegates to inline_graph.py."""
    user_codex_token = user.get("settings", {}).get("codex_token")
    user_openai_key = user.get("settings", {}).get("openai_key")

    return await run_inline_agent(
        selected_text=req.selected_text,
        action=req.action,
        context=req.context,
        model=req.model,
        codex_token=user_codex_token,
        openai_key=user_openai_key,
    )


class SaveEditRequest(BaseModel):
    conversation_id: str
    parent_report_id: int
    new_html_content: str


@router.post("/api/report-save-edit")
async def save_edited_report(req: SaveEditRequest, user: dict = Depends(get_current_user)):
    """Saves an inline-edited HTML report as a new version."""
    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()

        cur.execute("SELECT title, prompt, version FROM dashboard_canifa.saved_reports WHERE id = %s",
                     (req.parent_report_id,))
        row = cur.fetchone()
        if not row:
            from fastapi import HTTPException
            raise HTTPException(404, "Parent report not found")

        title, prompt, p_version = row[0], row[1], row[2]
        new_version = p_version + 1

        cur.execute(
            """INSERT INTO dashboard_canifa.saved_reports
            (title, prompt, html_content, tools_used, cycles_count, generation_time_ms, version, parent_id, conversation_id)
            VALUES (%s, %s, %s, '[]'::jsonb, 0, 0, %s, %s, %s)
            RETURNING id""",
            (title, prompt, req.new_html_content, new_version, req.parent_report_id, req.conversation_id),
        )
        new_id = cur.fetchone()[0]

        cur.execute(
            """INSERT INTO dashboard_canifa.chat_history
               (user_id, module, conversation_id, message, is_human)
               VALUES (1, 'report_generator', %s, %s, false)""",
            (req.conversation_id,
             f"Đã cập nhật báo cáo (Chỉnh sửa HTML): {title} (v{new_version})|REPORT:{new_id}")
        )

        cur.close()
        return {"success": True, "new_report_id": new_id, "version": new_version}
    except Exception as e:
        logger.error("Failed to save edited report to DB: %s", e)
        from fastapi import HTTPException
        raise HTTPException(500, str(e))
    finally:
        if conn: conn.close()


# ─── Follow-up API (SSE Streaming) ──────────────────────────────────

class FollowupRequest(BaseModel):
    question: str
    parent_report_id: int
    appendix_index: int = 0
    model: str = "codex/gpt-5.3-codex"


@router.post("/api/report-followup")
async def followup_stream(req: FollowupRequest, user: dict = Depends(get_current_user)):
    """SSE stream for follow-up sections (Phụ lục A, B, C...)."""
    user_codex_token = user.get("settings", {}).get("codex_token")
    user_openai_key = user.get("settings", {}).get("openai_key")

    async def event_generator():
        try:
            async for event in run_followup_agent(
                question=req.question,
                parent_report_id=req.parent_report_id,
                appendix_index=req.appendix_index,
                model=req.model,
                codex_token=user_codex_token,
                openai_key=user_openai_key,
            ):
                yield _sse(event)
        except Exception as e:
            logger.error("Follow-up stream error: %s", traceback.format_exc())
            yield _sse({"type": "error", "message": humanize_llm_error(e)})

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

