"""
Real-time Agent Session Viewer — WebSocket consumer + CLI tail command.

Architecture
------------
Each ``execute_agent_task`` sends **progress events** to a Redis pub/sub
channel (``agent:progress:<pipeline_run_id>``).  The WebSocket consumer
subscribes to this channel and streams events to connected clients.

A **lightweight CLI tail command** (`tail_agent`) also reads the same
stream and prints events to stdout — no browser needed.

Event Types
-----------
    agent.started        Agent began execution (model, prompt length)
    agent.thinking       Agent is assembling context / loading upstream
    agent.calling_llm    Agent is sending prompt to LLM
    agent.streaming      LLM tokens streaming (optional, chunked)
    agent.completed      Agent finished successfully
    agent.failed         Agent failed with error
    agent.retrying       Agent is retrying after failure
    pipeline.paused      Pipeline paused at manual gate
    pipeline.resumed     Pipeline resumed after gate approval

Channel Format
--------------
    agent:progress:<pipeline_run_id>
    agent:progress:all   (global stream — all pipelines)

Usage
-----
    # CLI (terminal-based real-time viewer)
    python manage.py tail_agent <pipeline_run_id>

    # WebSocket (connect from browser or wscat)
    wscat -c ws://localhost:8000/ws/agents/<pipeline_run_id>/

    # Global stream (all pipelines)
    python manage.py tail_agent --all
"""

from __future__ import annotations

import json
import time
import uuid
from datetime import datetime, timezone

import structlog

_log = structlog.get_logger("simorgh.agent_stream")


# ---------------------------------------------------------------------------
# Event emitter — called from execute_agent_task
# ---------------------------------------------------------------------------

def emit_agent_event(
    pipeline_run_id: str,
    agent_id: str,
    event_type: str,
    **kwargs,
) -> dict:
    """Emit a progress event to the agent stream.

    This function is called synchronously from within Celery tasks.
    It publishes to Redis pub/sub AND appends to the database log.

    Args:
        pipeline_run_id: PipelineRun public_id.
        agent_id: Two-digit agent ID, e.g. '07'.
        event_type: One of the event types above.
        **kwargs: Additional event data (model, prompt_length, error, etc.)

    Returns:
        The event dict that was emitted.
    """
    event = {
        "event": event_type,
        "pipeline_run_id": pipeline_run_id,
        "agent_id": agent_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "data": kwargs,
    }

    # 1. Publish to Redis pub/sub (if available)
    _publish_to_redis(pipeline_run_id, event)

    # 2. Log to structured logger
    _log.info(
        f"agent.progress.{event_type}",
        pipeline_run_id=pipeline_run_id,
        agent_id=agent_id,
        **{k: str(v)[:200] for k, v in kwargs.items()},
    )

    return event


def _publish_to_redis(pipeline_run_id: str, event: dict) -> None:
    """Publish event to Redis pub/sub channels."""
    try:
        from django.conf import settings

        redis_url = getattr(settings, "REDIS_URL", None)
        if not redis_url:
            return  # Redis not configured — skip pub/sub

        import redis as redis_lib

        r = redis_lib.from_url(redis_url)
        payload = json.dumps(event, ensure_ascii=False)

        # Per-pipeline channel
        r.publish(f"agent:progress:{pipeline_run_id}", payload)
        # Global channel
        r.publish("agent:progress:all", payload)

    except Exception:
        # Redis pub/sub is best-effort — never crash the task
        pass


# ---------------------------------------------------------------------------
# CLI event printer — used by tail_agent management command
# ---------------------------------------------------------------------------

class AgentSessionPrinter:
    """Formats agent progress events for CLI output.

    Mimics the look and feel of a Copilot chat window in the terminal.
    """

    STATUS_ICONS = {
        "agent.started":      "🚀",
        "agent.thinking":     "🧠",
        "agent.calling_llm":  "🤖",
        "agent.streaming":    "📝",
        "agent.completed":    "✅",
        "agent.failed":       "❌",
        "agent.retrying":     "🔄",
        "pipeline.paused":    "🛑",
        "pipeline.resumed":   "▶️",
    }

    AGENT_NAMES = {
        "00": "Module Planner",
        "01": "Domain Architect",
        "02": "Data Architect",
        "03": "Platform Integration",
        "04": "Workflow Architect",
        "05": "Event Architect",
        "06": "Search & AI",
        "07": "Backend Architect",
        "08": "Frontend & UX",
        "09": "Localization",
        "10": "Seed & Fixture",
        "11": "Security",
        "12": "Reporting & Analytics",
        "13": "Integration Architect",
        "14": "Performance",
        "15": "QA",
        "16": "Documentation",
        "17": "Architecture Review",
        "18": "DevOps & Production",
        "19": "CTO Review",
    }

    def __init__(self, pipeline_run_id: str | None = None):
        self.pipeline_run_id = pipeline_run_id
        self._agent_sessions: dict[str, dict] = {}  # agent_id → session state
        self._event_count = 0

    def print_event(self, event: dict) -> None:
        """Format and print a single event."""
        self._event_count += 1
        event_type = event["event"]
        agent_id = event.get("agent_id", "??")
        data = event.get("data", {})
        icon = self.STATUS_ICONS.get(event_type, "•")

        agent_name = self.AGENT_NAMES.get(agent_id, f"Agent {agent_id}")
        ts = datetime.fromisoformat(event["timestamp"]).strftime("%H:%M:%S")

        if event_type == "agent.started":
            model = data.get("model", "?")
            prompt_len = data.get("prompt_length", 0)
            print(f"\n{'='*70}")
            print(f"  {icon} Agent {agent_id} — {agent_name}")
            print(f"     Model: {model}  |  Prompt: {prompt_len:,} chars  |  {ts}")
            print(f"  {'─'*66}")
            self._agent_sessions[agent_id] = {"started_at": ts, "model": model}

        elif event_type == "agent.thinking":
            context = data.get("context", "")
            print(f"  {icon} Thinking... {context}")

        elif event_type == "agent.calling_llm":
            model = data.get("model", "")
            prompt_len = data.get("prompt_length", 0)
            print(f"  {icon} Calling LLM [{model}] — {prompt_len:,} chars")
            print(f"     (This may take 30-60 seconds for large agents...)")

        elif event_type == "agent.completed":
            output_len = data.get("output_length", 0)
            duration = data.get("duration_seconds", 0)
            output_file = data.get("output_file", "")
            mode = data.get("mode", "?")
            print(f"  {icon} Done!  Output: {output_len:,} chars  |  {duration:.1f}s  |  Mode: {mode}")
            if output_file:
                print(f"     📄 {output_file}")
            print(f"  {'─'*66}")

        elif event_type == "agent.failed":
            error = data.get("error", "")[:200]
            retry = data.get("retry", "")
            print(f"  {icon} FAILED {retry}")
            print(f"     Error: {error}")
            print(f"  {'─'*66}")

        elif event_type == "agent.retrying":
            attempt = data.get("attempt", "?")
            delay = data.get("delay_seconds", 0)
            print(f"  {icon} Retrying... attempt {attempt} in {delay}s")

        elif event_type == "pipeline.paused":
            after_agent = data.get("after_agent", "??")
            gate_desc = data.get("gate_description", "")
            print(f"\n{'='*70}")
            print(f"  {icon} PIPELINE PAUSED — After Agent {after_agent}")
            print(f"     {gate_desc}")
            print(f"     To approve: manage.py approve_gate {self.pipeline_run_id} --gate {after_agent}")
            print(f"  {'='*70}")

        elif event_type == "pipeline.resumed":
            print(f"\n  {icon} Pipeline RESUMED — continuing with next agents...\n")

        elif event_type == "agent.streaming":
            chunk = data.get("chunk", "")
            if chunk:
                print(chunk, end="", flush=True)


# ---------------------------------------------------------------------------
# WebSocket Consumer (Django Channels)
# ---------------------------------------------------------------------------
# If Django Channels is available, this consumer streams agent progress
# to browser clients in real-time.  The frontend can render these events
# in a Copilot-like chat panel.

try:
    from channels.generic.websocket import AsyncWebsocketConsumer

    class AgentProgressConsumer(AsyncWebsocketConsumer):
        """WebSocket consumer for agent progress streaming.

        Connect to: ws://<host>/ws/agents/<pipeline_run_id>/
        Or global:  ws://<host>/ws/agents/all/
        """

        async def connect(self):
            self.pipeline_run_id = self.scope["url_route"]["kwargs"].get("pipeline_run_id", "all")
            self.room_group_name = f"agent_progress_{self.pipeline_run_id}"

            await self.channel_layer.group_add(
                self.room_group_name,
                self.channel_name,
            )
            await self.accept()

            await self.send(text_data=json.dumps({
                "event": "connected",
                "pipeline_run_id": self.pipeline_run_id,
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }))

        async def disconnect(self, close_code):
            await self.channel_layer.group_discard(
                self.room_group_name,
                self.channel_name,
            )

        async def agent_event(self, event):
            """Receive event from Redis→Channels bridge and forward to WebSocket."""
            await self.send(text_data=json.dumps(event["data"]))

except ImportError:
    # Django Channels not installed — WebSocket consumer unavailable
    AgentProgressConsumer = None  # type: ignore