"""
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 fastapi import APIRouter, Depends
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from starlette.responses import StreamingResponse

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

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 get_task_status, read_events

    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")

