"""Main daemon event loop — one polling task per agent."""

import asyncio
import signal
import time
from datetime import datetime
from typing import Dict, Optional

import httpx

from mercury_launcher.agent_runner import AgentRunner
from mercury_launcher.context_builder import build_session_prompt, build_system_prompt
from mercury_launcher.log import get_logger
from mercury_launcher.models import AgentConfig, LauncherConfig
from mercury_launcher.poller import MercuryError, MercuryPoller

logger = get_logger("daemon")

# ── Rich terminal output helpers ──────────────────────────────────────────

# ANSI color palette for agent names (cycle through these)
_AGENT_COLORS = [
    "\033[36m",   # cyan
    "\033[33m",   # yellow
    "\033[35m",   # magenta
    "\033[32m",   # green
    "\033[34m",   # blue
    "\033[91m",   # bright red
]
_RESET = "\033[0m"
_DIM = "\033[2m"
_BOLD = "\033[1m"
_agent_color_map: Dict[str, str] = {}


def _color_for(agent_name: str) -> str:
    if agent_name not in _agent_color_map:
        idx = len(_agent_color_map) % len(_AGENT_COLORS)
        _agent_color_map[agent_name] = _AGENT_COLORS[idx]
    return _agent_color_map[agent_name]


def _print_activity(agent_name: str, text: str, dim: bool = False) -> None:
    """Print a timestamped, color-coded activity line."""
    ts = datetime.now().strftime("%H:%M:%S")
    color = _color_for(agent_name)
    style = _DIM if dim else ""
    print(f"{_DIM}{ts}{_RESET} {color}{_BOLD}{agent_name}{_RESET} {style}{text}{_RESET}")


class Daemon:
    """Manages parallel polling loops for all configured agents."""

    def __init__(self, config: LauncherConfig, agents: list[AgentConfig] | None = None):
        self.config = config
        self._agents = agents or config.agents
        self._shutdown = asyncio.Event()
        self._agent_states: Dict[str, str] = {}  # name -> state

    async def run(self) -> None:
        """Start all agent loops and block until shutdown."""
        if not self._agents:
            logger.error("No agents configured. Add agents to launcher.yaml.")
            return

        # Register signal handlers
        loop = asyncio.get_running_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            loop.add_signal_handler(sig, self._handle_signal)

        print(f"\n{_BOLD}Mercury Launcher{_RESET} — {len(self._agents)} agent(s)\n")
        for agent in self._agents:
            color = _color_for(agent.name)
            print(f"  {color}{_BOLD}{agent.name}{_RESET} → {_DIM}{agent.working_directory}{_RESET}")
        print()

        tasks = [
            asyncio.create_task(self._agent_loop(agent), name=f"agent:{agent.name}")
            for agent in self._agents
        ]

        # Wait for shutdown signal
        await self._shutdown.wait()
        logger.info("Shutting down...")

        # Cancel all agent loops
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)

        logger.info("All agent loops stopped.")

    def _handle_signal(self) -> None:
        logger.info("Received shutdown signal")
        self._shutdown.set()

    async def _agent_loop(self, agent_config: AgentConfig) -> None:
        """Poll-and-spawn loop for a single agent."""
        name = agent_config.name
        self._agent_states[name] = "starting"

        poller = MercuryPoller(
            endpoint=agent_config.mercury_url,
            api_key=agent_config.mcp_api_key,
        )

        async def _on_status(status: str) -> None:
            try:
                await poller.update_status(status)
            except Exception:
                pass

        # Activity posting state
        _active_thread_ids: list[str] = []
        _last_activity_post: Dict[str, float] = {}  # thread_id -> time.time()
        _ACTIVITY_THROTTLE_S = 5.0

        async def _on_activity(label: str, text_content: Optional[str]) -> None:
            """Rich terminal output + throttled chat activity posting."""
            if text_content:
                _print_activity(name, text_content, dim=True)
            else:
                _print_activity(name, label)
            # Post to active threads (throttled)
            now = time.time()
            for tid in _active_thread_ids:
                last = _last_activity_post.get(tid, 0)
                if now - last < _ACTIVITY_THROTTLE_S:
                    continue
                _last_activity_post[tid] = now
                try:
                    await poller.post_activity(tid, label)
                except Exception:
                    pass

        runner = AgentRunner(
            agent_name=name,
            mercury_url=agent_config.mercury_url,
            mcp_api_key=agent_config.mcp_api_key,
            working_directory=agent_config.working_directory,
            model=agent_config.model,
            max_budget_usd=agent_config.max_budget_usd,
            on_status=_on_status,
            on_activity=_on_activity,
        )

        # Fetch agent context once at startup (cached for session prompts)
        agent_context = {}
        system_prompt = ""

        try:
            agent_context = await poller.get_agent_context()
            system_prompt = build_system_prompt(agent_context)
            ctx_name = agent_context.get("name", name)
            ctx_role = agent_context.get("agent_role", "")
            _print_activity(name, f"Context loaded — {ctx_role}" if ctx_role else "Context loaded")
        except MercuryError as e:
            logger.warning("[%s] Could not fetch agent context: %s", name, e)
            system_prompt = f"You are {name}, a Mercury agent. Respond to messages using Mercury MCP tools."
        except Exception as e:
            logger.warning("[%s] Error fetching agent context: %s", name, e)
            system_prompt = f"You are {name}, a Mercury agent. Respond to messages using Mercury MCP tools."

        self._agent_states[name] = "polling"
        _print_activity(name, "Listening for messages...")

        while not self._shutdown.is_set():
            try:
                # Long-poll for messages (blocks up to 30s)
                messages = await poller.wait_for_messages(timeout=30)

                if not messages:
                    continue

                # Track active thread IDs for activity posting
                _active_thread_ids.clear()
                _last_activity_post.clear()
                for msg in messages:
                    sender = msg.get("sender_name", "?")
                    preview = msg.get("content", "")[:100]
                    _print_activity(name, f"Message from {_BOLD}{sender}{_RESET}: {preview}")
                    tid = msg.get("thread_id", "")
                    if tid and tid not in _active_thread_ids:
                        _active_thread_ids.append(tid)

                # Build session prompt from incoming messages
                session_prompt = build_session_prompt(messages)

                # Spawn Claude Code session
                self._agent_states[name] = "session_active"
                try:
                    result = await runner.run_session(system_prompt, session_prompt)
                finally:
                    # Clear status in Mercury UI after session ends
                    try:
                        await poller.update_status(None)
                    except Exception:
                        pass

                # Post session summary
                cost = result.get("cost_usd", 0)
                turns = result.get("num_turns", 0)
                if result.get("status") == "error":
                    _print_activity(name, f"\033[91mSession failed: {result.get('error')}{_RESET}")
                else:
                    _print_activity(name, f"Session complete — {turns} turns, ${cost:.4f}")
                    summary = f"Session complete — {turns} turns, ${cost:.4f}"
                    for tid in _active_thread_ids:
                        try:
                            await poller.post_activity(tid, summary)
                        except Exception:
                            pass
                _active_thread_ids.clear()

                self._agent_states[name] = "polling"

                # Cooldown between sessions
                await asyncio.sleep(agent_config.cooldown_seconds)

            except httpx.ConnectError:
                logger.warning("[%s] Connection refused — is Mercury running? Retrying in 10s…", name)
                self._agent_states[name] = "reconnecting"
                await asyncio.sleep(10)

            except httpx.HTTPStatusError as e:
                status = e.response.status_code
                if status in (502, 503, 504):
                    logger.warning("[%s] Server temporarily unavailable (%d). Retrying in 10s…", name, status)
                    self._agent_states[name] = "reconnecting"
                    await asyncio.sleep(10)
                else:
                    logger.error("[%s] HTTP error %d: %s", name, status, e)
                    await asyncio.sleep(5)

            except MercuryError as e:
                logger.error("[%s] Mercury error: %s", name, e)
                await asyncio.sleep(5)

            except asyncio.CancelledError:
                break

            except Exception as e:
                logger.error("[%s] Unexpected error: %s", name, e, exc_info=True)
                await asyncio.sleep(5)

        # Cleanup — clear status before closing
        try:
            await poller.update_status(None)
        except Exception:
            pass
        await poller.close()
        self._agent_states[name] = "stopped"
        logger.info("[%s] Stopped", name)

    def get_status(self) -> Dict[str, str]:
        """Return current state of all agents."""
        return dict(self._agent_states)
