"""
Real-time agent session viewer — CLI-based, like a Copilot chat window.

Usage
-----
.. code-block:: sh

    # Watch a specific pipeline
    python manage.py tail_agent <pipeline_run_id>

    # Watch ALL pipelines (global stream)
    python manage.py tail_agent --all

    # Watch with verbose mode (full event JSON)
    python manage.py tail_agent <pipeline_run_id> --verbose

    # Watch for a specific agent only
    python manage.py tail_agent <pipeline_run_id> --agent 07

How it works
------------
Subscribes to Redis pub/sub channels and prints events in real-time.
If Redis is not available, polls the database AgentRun table every 2 seconds.

Output looks like:
    ======================================================================
      🚀 Agent 02 — Data Architect
         Model: deepseek-v4-pro  |  Prompt: 12,345 chars  |  14:30:05
      ──────────────────────────────────────────────────────────────────
      🧠 Thinking... Loading upstream outputs...
      🤖 Calling LLM [deepseek-v4-pro] — 12,345 chars
         (This may take 30-60 seconds for large agents...)
      ✅ Done!  Output: 8,450 chars  |  42.3s  |  Mode: llm
         📄 .docs/generated/inbox_referral/02-data-model.md
      ──────────────────────────────────────────────────────────────────
    ======================================================================
      🛑 PIPELINE PAUSED — After Agent 02
         مدل داده — بررسی، ویرایش و تأیید
         To approve: manage.py approve_gate abc-123 --gate 02
    ======================================================================
"""

from __future__ import annotations

import json
import time
from pathlib import Path

from django.core.management.base import BaseCommand, CommandError


class Command(BaseCommand):
    help = __doc__

    def add_arguments(self, parser):
        parser.add_argument(
            "pipeline_run_id",
            nargs="?",
            help="PipelineRun public_id (UUID). Omit with --all.",
        )
        parser.add_argument(
            "--all",
            action="store_true",
            help="Watch ALL pipelines (global stream).",
        )
        parser.add_argument(
            "--agent",
            help="Filter events to a specific agent ID, e.g. '07'.",
        )
        parser.add_argument(
            "--verbose",
            action="store_true",
            help="Show full event JSON instead of formatted output.",
        )
        parser.add_argument(
            "--poll",
            action="store_true",
            help="Force database polling mode (even if Redis is available).",
        )

    def handle(self, **options):
        if not options["pipeline_run_id"] and not options["all"]:
            raise CommandError(
                "Provide a pipeline_run_id or use --all for global stream.\n"
                "Usage: manage.py tail_agent <pipeline_run_id>"
            )

        pipeline_run_id = options["pipeline_run_id"] if not options["all"] else None
        agent_filter = options.get("agent")

        if options["verbose"]:
            self._tail_verbose(pipeline_run_id, agent_filter)
            return

        if options["poll"] or not self._redis_available():
            self._tail_poll(pipeline_run_id, agent_filter)
        else:
            self._tail_redis(pipeline_run_id, agent_filter)

    # ------------------------------------------------------------------
    # Redis pub/sub mode (real-time)
    # ------------------------------------------------------------------

    def _redis_available(self) -> bool:
        """Check if Redis is configured."""
        try:
            from django.conf import settings
            return bool(getattr(settings, "REDIS_URL", None))
        except Exception:
            return False

    def _tail_redis(self, pipeline_run_id: str | None, agent_filter: str | None) -> None:
        """Subscribe to Redis pub/sub and print events."""
        from django.conf import settings
        import redis as redis_lib
        from simorgh.apps.automation.agent_stream import AgentSessionPrinter

        r = redis_lib.from_url(settings.REDIS_URL)
        pubsub = r.pubsub()

        channel = "agent:progress:all" if pipeline_run_id is None else f"agent:progress:{pipeline_run_id}"
        pubsub.subscribe(channel)

        printer = AgentSessionPrinter(pipeline_run_id or "all")

        self.stdout.write(self.style.SUCCESS(
            f"\n🔍 Watching agent sessions... (Ctrl+C to stop)\n"
            f"   Channel: {channel}\n"
            f"   Filter:  {'all agents' if not agent_filter else f'agent {agent_filter}'}\n"
        ))

        try:
            for message in pubsub.listen():
                if message["type"] != "message":
                    continue
                try:
                    event = json.loads(message["data"])
                    if agent_filter and event.get("agent_id") != agent_filter:
                        continue
                    printer.print_event(event)
                except (json.JSONDecodeError, KeyError):
                    pass
        except KeyboardInterrupt:
            self.stdout.write("\n👋 Disconnected.")

    # ------------------------------------------------------------------
    # Database polling mode (fallback when Redis unavailable)
    # ------------------------------------------------------------------

    def _tail_poll(self, pipeline_run_id: str | None, agent_filter: str | None) -> None:
        """Poll the database AgentRun table for changes."""
        from simorgh.apps.automation.agent_models import AgentRun, AgentRunStatus, PipelineRun
        from simorgh.apps.automation.agent_stream import AgentSessionPrinter

        if pipeline_run_id:
            try:
                pr = PipelineRun.objects.get(public_id=pipeline_run_id)
            except PipelineRun.DoesNotExist:
                raise CommandError(f"PipelineRun not found: {pipeline_run_id}")
            runs = [pr]
        else:
            runs = list(PipelineRun.objects.filter(status__in=["running", "paused"]).order_by("-created_at")[:5])

        if not runs:
            self.stdout.write(self.style.WARNING("No active pipelines found."))
            return

        printer = AgentSessionPrinter(pipeline_run_id or "all")
        seen_states: dict[str, dict] = {}  # key: "{pipeline_id}:{agent_id}" → last status

        self.stdout.write(self.style.SUCCESS(
            f"\n🔍 Polling agent sessions every 2s... (Ctrl+C to stop)\n"
            f"   Mode: database polling (Redis not available)\n"
            f"   Pipelines: {len(runs)}\n"
        ))

        try:
            while True:
                for pr in runs:
                    pr.refresh_from_db()
                    for ar in pr.agent_runs.all():
                        key = f"{pr.public_id}:{ar.agent_id}"
                        last = seen_states.get(key, {})

                        if last.get("status") != ar.status:
                            event = self._agent_run_to_event(pr, ar)
                            if agent_filter and event.get("agent_id") != agent_filter:
                                continue
                            printer.print_event(event)
                            seen_states[key] = {
                                "status": ar.status,
                                "updated_at": str(ar.updated_at),
                            }

                    if pr.status == "paused" and seen_states.get(f"{pr.public_id}:__paused__") != pr.status:
                        printer.print_event({
                            "event": "pipeline.paused",
                            "pipeline_run_id": str(pr.public_id),
                            "agent_id": "??",
                            "timestamp": str(pr.updated_at),
                            "data": {
                                "after_agent": str(pr.current_step),
                                "gate_description": "Pipeline paused — awaiting manual gate approval.",
                            },
                        })
                        seen_states[f"{pr.public_id}:__paused__"] = pr.status

                time.sleep(2)

        except KeyboardInterrupt:
            self.stdout.write("\n👋 Disconnected.")

    def _agent_run_to_event(self, pr, ar) -> dict:
        """Convert an AgentRun DB row to a progress event."""
        from datetime import timezone as tz

        status_event_map = {
            "running": "agent.started",
            "completed": "agent.completed",
            "failed": "agent.failed",
            "pending": "agent.started",
        }

        return {
            "event": status_event_map.get(ar.status, "agent.started"),
            "pipeline_run_id": str(pr.public_id),
            "agent_id": ar.agent_id,
            "timestamp": ar.updated_at.replace(tzinfo=tz.utc).isoformat() if ar.updated_at else "",
            "data": {
                "model": ar.model_name or pr.default_model,
                "output_length": ar.output_summary.get("output_length", 0) if ar.output_summary else 0,
                "output_file": ar.output_summary.get("file_path", "") if ar.output_summary else "",
                "mode": ar.output_summary.get("mode", "?") if ar.output_summary else "?",
                "error": ar.error_message[:200] if ar.error_message else "",
                "retry": f"Retry {ar.retry_count}" if ar.retry_count else "",
            },
        }

    # ------------------------------------------------------------------
    # Verbose mode (full JSON)
    # ------------------------------------------------------------------

    def _tail_verbose(self, pipeline_run_id: str | None, agent_filter: str | None) -> None:
        """Print raw JSON events — same as Redis mode but formatted JSON."""
        try:
            from django.conf import settings
            import redis as redis_lib

            r = redis_lib.from_url(settings.REDIS_URL)
            pubsub = r.pubsub()
            channel = "agent:progress:all" if pipeline_run_id is None else f"agent:progress:{pipeline_run_id}"
            pubsub.subscribe(channel)

            self.stdout.write(f"📡 Subscribed to {channel} (Ctrl+C to stop)\n")

            for message in pubsub.listen():
                if message["type"] != "message":
                    continue
                try:
                    event = json.loads(message["data"])
                    if agent_filter and event.get("agent_id") != agent_filter:
                        continue
                    self.stdout.write(json.dumps(event, indent=2, ensure_ascii=False))
                    self.stdout.write("─" * 60)
                except json.JSONDecodeError:
                    pass
        except KeyboardInterrupt:
            self.stdout.write("\n👋 Disconnected.")
        except ImportError:
            self.stdout.write(self.style.ERROR("Redis not available. Use --poll for database mode."))