"""
Phase 19: Swarm Agent Discovery

Dynamic agent registry for the swarm. Agents advertise capabilities,
declare availability, and discover peers for task delegation.
Includes deduplication, staleness cleanup, fuzzy capability matching,
and load-aware agent selection.
"""

import json
import time
from pathlib import Path
from typing import Optional


STALE_THRESHOLD_SEC = 600  # 10 minutes without heartbeat = stale


class AgentDiscovery:
    def __init__(self, storage_dir: str = "state/broker/discovery"):
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self.registry_file = self.storage_dir / "agent_registry.json"
        if not self.registry_file.exists():
            self._save({})

    def _load(self) -> dict:
        try:
            with open(self.registry_file, "r", encoding="utf-8") as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError):
            return {}

    def _save(self, data: dict) -> None:
        with open(self.registry_file, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)

    # ── Registration ──────────────────────────────────────────────────

    def register_agent(
        self,
        agent_id: str,
        capabilities: list[str],
        platform: str = "",
        display_name: str = "",
        max_concurrent_tasks: int = 3,
        hardware: Optional[dict] = None,
    ) -> dict:
        """Register or update an agent in the discovery registry.

        If agent_id already exists, updates their info (no duplicates).
        Returns the agent's registry entry.
        """
        data = self._load()

        # Normalize capabilities (lowercase, dedup, sorted)
        caps = sorted(set(c.lower().strip() for c in capabilities if c.strip()))

        existing = data.get(agent_id, {})
        data[agent_id] = {
            "agent_id": agent_id,
            "display_name": display_name or existing.get("display_name", agent_id),
            "platform": platform or existing.get("platform", "unknown"),
            "capabilities": caps,
            "max_concurrent_tasks": max_concurrent_tasks,
            "hardware": hardware or existing.get("hardware", {}),
            "last_seen": time.time(),
            "registered_at": existing.get("registered_at", time.time()),
            "status": "available",
            "current_tasks": existing.get("current_tasks", 0),
        }
        self._save(data)
        return data[agent_id]

    def unregister_agent(self, agent_id: str) -> dict:
        """Remove an agent from the registry."""
        data = self._load()
        if agent_id not in data:
            return {"ok": False, "error": "Agent not found"}
        del data[agent_id]
        self._save(data)
        return {"ok": True, "removed": agent_id}

    # ── Heartbeat ─────────────────────────────────────────────────────

    def heartbeat(self, agent_id: str, current_tasks: int = 0) -> dict:
        """Update agent's heartbeat. Auto-registers if unknown."""
        data = self._load()
        if agent_id not in data:
            # Auto-register with empty capabilities — agent should call register_agent next
            data[agent_id] = {
                "agent_id": agent_id,
                "display_name": agent_id,
                "platform": "unknown",
                "capabilities": [],
                "max_concurrent_tasks": 3,
                "hardware": {},
                "last_seen": time.time(),
                "registered_at": time.time(),
                "status": "available",
                "current_tasks": current_tasks,
            }
        else:
            data[agent_id]["last_seen"] = time.time()
            data[agent_id]["current_tasks"] = current_tasks
            # Restore to available if was stale
            if data[agent_id].get("status") == "stale":
                data[agent_id]["status"] = "available"
        self._save(data)
        return {"ok": True, "agent_id": agent_id}

    # ── Discovery ─────────────────────────────────────────────────────

    def find_agents_by_capability(
        self,
        capability: str,
        include_stale: bool = False,
        sort_by_load: bool = True,
    ) -> list[dict]:
        """Find all agents with a given capability.

        Supports exact and substring matching. Excludes stale agents
        by default. Optionally sorts by load (least busy first).
        """
        data = self._load()
        self._mark_stale(data)
        capability_lower = capability.lower().strip()
        matched = []

        for agent_id, info in data.items():
            if not include_stale and info.get("status") == "stale":
                continue

            # Exact or substring match on capabilities
            caps = info.get("capabilities", [])
            match = any(capability_lower == c or capability_lower in c or c in capability_lower for c in caps)
            if match:
                info_copy = {**info}
                # Score: exact match > substring
                if capability_lower in caps:
                    info_copy["_match_type"] = "exact"
                    info_copy["_match_score"] = 100
                else:
                    info_copy["_match_type"] = "partial"
                    info_copy["_match_score"] = 60
                matched.append(info_copy)

        if sort_by_load:
            # Least busy first, then highest match score
            matched.sort(key=lambda x: (x.get("current_tasks", 0), -x.get("_match_score", 0)))
        else:
            matched.sort(key=lambda x: -x.get("_match_score", 0))

        return matched

    def find_agents_by_hardware(self, requirement: str) -> list[dict]:
        """Find agents that have declared specific hardware (e.g., 'gpu', 'browser')."""
        data = self._load()
        self._mark_stale(data)
        requirement_lower = requirement.lower().strip()
        matched = []

        for agent_id, info in data.items():
            if info.get("status") == "stale":
                continue
            hw = info.get("hardware", {})
            # Check if requirement key exists and is truthy
            if hw.get(requirement_lower) or any(requirement_lower in k.lower() for k in hw):
                matched.append(info)

        return matched

    def list_agents(self, include_stale: bool = False) -> list[dict]:
        """List all registered agents."""
        data = self._load()
        self._mark_stale(data)
        agents = []
        for agent_id, info in data.items():
            if not include_stale and info.get("status") == "stale":
                continue
            agents.append(info)
        agents.sort(key=lambda x: x.get("last_seen", 0), reverse=True)
        return agents

    def get_agent(self, agent_id: str) -> Optional[dict]:
        """Get details for a specific agent."""
        data = self._load()
        return data.get(agent_id)

    # ── Staleness Management ──────────────────────────────────────────

    def _mark_stale(self, data: dict) -> int:
        """Mark agents as stale if they haven't sent a heartbeat recently.
        Modifies data in-place and saves. Returns count of newly stale agents."""
        now = time.time()
        newly_stale = 0
        for agent_id, info in data.items():
            last_seen = info.get("last_seen", 0)
            if (now - last_seen) > STALE_THRESHOLD_SEC and info.get("status") != "stale":
                info["status"] = "stale"
                newly_stale += 1
        if newly_stale > 0:
            self._save(data)
        return newly_stale

    def cleanup_stale(self, max_stale_hours: int = 24) -> dict:
        """Remove agents that have been stale for more than max_stale_hours."""
        data = self._load()
        now = time.time()
        cutoff = max_stale_hours * 3600
        removed = []

        for agent_id in list(data.keys()):
            info = data[agent_id]
            last_seen = info.get("last_seen", 0)
            if (now - last_seen) > cutoff:
                removed.append(agent_id)
                del data[agent_id]

        if removed:
            self._save(data)
        return {"removed": removed, "remaining": len(data)}

    # ── Stats ─────────────────────────────────────────────────────────

    def stats(self) -> dict:
        """Get discovery registry statistics."""
        data = self._load()
        self._mark_stale(data)
        total = len(data)
        active = sum(1 for v in data.values() if v.get("status") != "stale")
        stale = total - active
        all_caps = set()
        for info in data.values():
            all_caps.update(info.get("capabilities", []))
        return {
            "total_agents": total,
            "active_agents": active,
            "stale_agents": stale,
            "unique_capabilities": len(all_caps),
            "capabilities": sorted(all_caps),
        }
