"""
Report Queue — Redis-backed task queue for async report generation.

Uses Redis Streams for ordered event delivery + Hash for task metadata.
Worker runs as asyncio.create_task() inside the same uvicorn process.

Key structure:
  report:task:{task_id}     — Hash: status, question, model, etc.
  report:events:{task_id}   — Stream: ordered SSE events
  report:queue              — List: pending task_ids (FIFO)
"""

import asyncio
import json
import logging
import time
import uuid

import psycopg

from common.cache import redis_cache
from config import CHECKPOINT_POSTGRES_URL
from common.pool_wrapper import get_pooled_connection_compat
from agent.report_agent.main_graph import run_report_agent
from agent.report_agent.core import humanize_llm_error

logger = logging.getLogger(__name__)

QUEUE_KEY = "report:queue"
TASK_PREFIX = "report:task:"
EVENTS_PREFIX = "report:events:"
TASK_TTL = 7200  # 2 hours


# ─── Task Management ────────────────────────────────────────────────

async def enqueue_report_task(
    *,
    question: str,
    model: str,
    conversation_id: str,
    parent_report_id: int | None = None,
    current_html: str | None = None,
    codex_token: str | None = None,
    openai_key: str | None = None,
) -> str | None:
    """
    Enqueue a report generation task into Redis.
    Returns task_id or None if Redis unavailable (fallback to direct SSE).
    """
    client = redis_cache.get_client()
    if not client:
        return None

    task_id = str(uuid.uuid4())
    task_key = f"{TASK_PREFIX}{task_id}"

    try:
        task_data = {
            "status": "pending",
            "question": question,
            "model": model,
            "conversation_id": conversation_id,
            "parent_report_id": str(parent_report_id) if parent_report_id else "",
            "current_html": current_html or "",
            "codex_token": codex_token or "",
            "openai_key": openai_key or "",
            "created_at": str(time.time()),
        }
        await client.hset(task_key, mapping=task_data)
        await client.expire(task_key, TASK_TTL)
        await client.lpush(QUEUE_KEY, task_id)
        logger.info("📥 Report task enqueued: %s — %s", task_id, question[:60])
        return task_id

    except Exception as e:
        logger.error("Failed to enqueue report task: %s", e)
        return None


async def get_task_status(task_id: str) -> dict | None:
    """Get task metadata from Redis."""
    client = redis_cache.get_client()
    if not client:
        return None
    try:
        return await client.hgetall(f"{TASK_PREFIX}{task_id}")
    except Exception:
        return None


async def write_event(task_id: str, event: dict):
    """Write a single event to the Redis Stream for this task."""
    client = redis_cache.get_client()
    if not client:
        return
    try:
        # Serialize all values to strings for Redis Stream
        flat = {k: json.dumps(v, ensure_ascii=False, default=str) if not isinstance(v, str) else v
                for k, v in event.items()}
        await client.xadd(
            f"{EVENTS_PREFIX}{task_id}",
            flat,
            maxlen=10000,
        )
    except Exception as e:
        logger.warning("Failed to write event to Redis: %s", e)


async def read_events(task_id: str, last_id: str = "0-0", block_ms: int = 2000, count: int = 50):
    """
    Read events from Redis Stream since last_id.
    Returns list of (msg_id, event_dict) tuples.
    """
    client = redis_cache.get_client()
    if not client:
        return []
    try:
        stream_key = f"{EVENTS_PREFIX}{task_id}"
        result = await client.xread({stream_key: last_id}, count=count, block=block_ms)
        if not result:
            return []

        events = []
        for stream_name, messages in result:
            for msg_id, data in messages:
                # Deserialize JSON values back
                parsed = {}
                for k, v in data.items():
                    try:
                        parsed[k] = json.loads(v)
                    except (json.JSONDecodeError, TypeError):
                        parsed[k] = v
                events.append((msg_id, parsed))
        return events

    except Exception as e:
        logger.warning("Failed to read events from Redis: %s", e)
        return []


# ─── Worker Loop ─────────────────────────────────────────────────────

async def report_worker_loop():
    """
    Background worker — runs forever, picks tasks from Redis queue.
    Called via asyncio.create_task() on FastAPI startup.
    """
    logger.info("🔄 Report Worker started — waiting for tasks...")

    while True:
        client = redis_cache.get_client()
        if not client:
            await asyncio.sleep(5)
            continue

        try:
            # Blocking pop with 5s timeout
            result = await client.brpop(QUEUE_KEY, timeout=5)
            if not result:
                continue

            _, task_id = result
            if isinstance(task_id, bytes):
                task_id = task_id.decode()

            task_key = f"{TASK_PREFIX}{task_id}"
            task_data = await client.hgetall(task_key)

            if not task_data or not task_data.get("question"):
                logger.warning("Empty task data for %s, skipping", task_id)
                continue

            logger.info("🚀 Worker processing: %s — %s", task_id, task_data["question"][:60])
            await client.hset(task_key, "status", "running")

            await _process_task(task_id, task_data)

        except asyncio.CancelledError:
            logger.info("🛑 Report Worker shutting down")
            break
        except Exception as e:
            logger.error("Worker loop error: %s", e)
            await asyncio.sleep(2)


async def _process_task(task_id: str, task_data: dict):
    """Process a single report task — run agent and write events to Redis."""
    client = redis_cache.get_client()
    task_key = f"{TASK_PREFIX}{task_id}"
    conv_id = task_data.get("conversation_id", "")
    start_time = time.time()

    try:
        # Save user message to chat_history
        _save_chat_msg(conv_id, task_data["question"], is_human=True)

        parent_id = int(task_data["parent_report_id"]) if task_data.get("parent_report_id") else None
        current_html = task_data.get("current_html") or None
        codex_token = task_data.get("codex_token") or None
        openai_key = task_data.get("openai_key") or None

        html_body = ""
        tools_used = []
        cycles_count = 0

        async for event in run_report_agent(
            question=task_data["question"],
            model=task_data.get("model", ""),
            codex_token=codex_token,
            openai_key=openai_key,
            parent_report_id=parent_id,
            current_html=current_html,
        ):
            event_type = event.get("type")

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

                # Save report to DB
                generation_time_ms = int((time.time() - start_time) * 1000)
                report_id = _save_report_sync(
                    question=task_data["question"],
                    html_body=html_body,
                    tools_used=tools_used,
                    cycles_count=cycles_count,
                    generation_time_ms=generation_time_ms,
                    parent_report_id=parent_id,
                    current_html=current_html,
                    conv_id=conv_id,
                )

                # Emit enriched html_complete with report_id
                await write_event(task_id, {
                    "type": "html_complete",
                    "html": html_body,
                    "report_id": report_id,
                    "conversation_id": conv_id,
                })
                continue

            if event_type == "direct_response":
                message = event.get("message", "")
                await write_event(task_id, event)
                _save_chat_msg(conv_id, message, is_human=False)
                continue

            # Pass all other events through
            await write_event(task_id, event)

        # Final done event
        await write_event(task_id, {"type": "done"})
        if client:
            await client.hset(task_key, "status", "done")
            await client.expire(f"{EVENTS_PREFIX}{task_id}", TASK_TTL)

        logger.info("✅ Task %s completed in %.1fs", task_id, time.time() - start_time)

    except Exception as e:
        friendly_error = humanize_llm_error(e)
        logger.error("Task %s failed: %s", task_id, e)
        await write_event(task_id, {"type": "error", "message": friendly_error[:500]})
        if client:
            await client.hset(task_key, "status", "error")


# ─── Sync DB Helpers (used by worker) ────────────────────────────────

def _save_chat_msg(conv_id: str, message: str, is_human: bool):
    """Save message to chat_history (sync)."""
    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, %s)""",
            (conv_id, message[:10000], is_human),
        )
        cur.close()
    except Exception as e:
        logger.warning("Failed to save chat msg: %s", e)
    finally:
        if conn:
            conn.close()


def _save_report_sync(
    *,
    question: str,
    html_body: str,
    tools_used: list,
    cycles_count: int,
    generation_time_ms: int,
    parent_report_id: int | None,
    current_html: str | None,
    conv_id: str,
) -> int:
    """Save report to DB (sync). Returns report_id."""
    version = 1
    title = question[:100] + ("..." if len(question) > 100 else "")
    final_html = html_body

    conn = None
    try:
        conn = get_pooled_connection_compat()
        cur = conn.cursor()

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

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

        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, %s::jsonb, %s, %s, %s, %s, %s)
            RETURNING id""",
            (title, question, final_html, json.dumps(tools_used), cycles_count,
             generation_time_ms, version, parent_report_id, conv_id),
        )
        report_id = cur.fetchone()[0]

        action_label = "tạo" if version == 1 else "cập nhật"
        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, f"Đã {action_label} báo cáo: {title} (v{version})|REPORT:{report_id}"),
        )
        cur.close()
        logger.info("💾 Report saved: id=%d, v%d", report_id, version)
        return report_id

    except Exception as e:
        logger.error("Failed to save report: %s", e)
        return -1
    finally:
        if conn:
            conn.close()
