"""
Capability Registry for Brynq v2 — Phase 5.

Agents advertise what they can do. The broker routes tasks to the right agent.
Supports fuzzy matching, load balancing, and auto-delegation.

Storage:
  state/broker/capabilities/{session_id}_{capability}.json — individual registrations
  state/broker/capabilities/_task_log.jsonl — delegation history for load balancing
"""

import json
import os
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from brynq.server import (
    BROKER_DIR,
    CHANNELS_DIR,
    SESSIONS_DIR,
    SESSION_ID,
    SESSION_NAME,
    PLATFORM,
    _append_jsonl_sync,
    _write_json_sync,
    _read_jsonl_sync,
)

CAPABILITY_DIR = BROKER_DIR / "capabilities"
CAPABILITY_DIR.mkdir(parents=True, exist_ok=True)
TASK_LOG = CAPABILITY_DIR / "_task_log.jsonl"


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


def register_capability(
    capability: str,
    description: str,
    session_id: Optional[str] = None,
    session_name: Optional[str] = None,
    platform: Optional[str] = None,
) -> dict:
    """Register a capability for a session."""
    sid = session_id or SESSION_ID
    sname = session_name or SESSION_NAME
    plat = platform or PLATFORM

    data = {
        "session_id": sid,
        "session_name": sname,
        "platform": plat,
        "capability": capability,
        "description": description,
        "registered_at": datetime.now(timezone.utc).isoformat(),
    }

    path = CAPABILITY_DIR / f"{sid}_{capability}.json"
    _write_json_sync(path, data)

    # Announce to #_system
    entry = {
        "id": uuid.uuid4().hex[:12],
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "session_id": sid,
        "session_name": sname,
        "platform": plat,
        "channel": "_system",
        "msg_type": "status",
        "message": f"{sname}@{plat} registered capability: {capability} — {description}",
    }
    _append_jsonl_sync(CHANNELS_DIR / "_system.jsonl", entry)
    return data


def unregister_capability(capability: str, session_id: Optional[str] = None) -> bool:
    """Remove a capability registration."""
    sid = session_id or SESSION_ID
    path = CAPABILITY_DIR / f"{sid}_{capability}.json"
    if path.exists():
        path.unlink(missing_ok=True)
        return True
    return False


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


def _get_active_session_ids() -> set[str]:
    """Get session IDs that have active session files."""
    active = set()
    if SESSIONS_DIR.exists():
        for f in SESSIONS_DIR.glob("*.json"):
            active.add(f.stem)
    return active


def list_capabilities(active_only: bool = True) -> list[dict]:
    """List all registered capabilities, optionally filtered to active sessions."""
    caps = []
    for f in sorted(CAPABILITY_DIR.glob("*.json")):
        if f.name.startswith("_"):
            continue
        try:
            caps.append(json.loads(f.read_text("utf-8")))
        except (json.JSONDecodeError, OSError):
            continue

    if active_only:
        active = _get_active_session_ids()
        caps = [c for c in caps if c.get("session_id") in active]

    return caps


def route_task(task_title: str, required_capability: str) -> Optional[str]:
    """Find the best active session that has the requested capability.
    This restores backward compatibility for the dashboard UI task router.
    """
    agents = find_agents(capability=required_capability)
    if not agents:
        return None
    # Just take the first one (highest score from find_agents)
    best_match = agents[0]
    return best_match.get("session_name")


def find_agents(
    capability: Optional[str] = None,
    query: Optional[str] = None,
    platform: Optional[str] = None,
) -> list[dict]:
    """Find agents matching criteria. Supports fuzzy matching on capability names
    and descriptions. Returns ranked results with delegation counts for load balancing."""
    all_caps = list_capabilities(active_only=True)
    task_counts = _get_delegation_counts()

    scored: list[tuple[float, dict]] = []
    for cap in all_caps:
        score = 0.0
        cap_name = cap.get("capability", "").lower()
        cap_desc = cap.get("description", "").lower()
        cap_platform = cap.get("platform", "").lower()

        # Exact capability match
        if capability:
            cap_query = capability.lower()
            if cap_name == cap_query:
                score += 100
            elif cap_query in cap_name:
                score += 60
            elif cap_query in cap_desc:
                score += 30

        # Fuzzy query match across name + description
        if query:
            q = query.lower()
            words = q.split()
            for word in words:
                if word in cap_name:
                    score += 40
                if word in cap_desc:
                    score += 20

        # Platform filter (exact)
        if platform:
            if cap_platform != platform.lower():
                continue

        # If no criteria given, include everything
        if not capability and not query:
            score = 50

        if score <= 0:
            continue

        # Load balancing: reduce score for overloaded agents
        sid = cap.get("session_id", "")
        active_tasks = task_counts.get(sid, 0)
        score -= active_tasks * 5  # Penalize busy agents

        cap_result = {
            **cap,
            "score": round(score, 1),
            "active_delegations": active_tasks,
        }
        scored.append((score, cap_result))

    # Sort by score descending
    scored.sort(key=lambda x: -x[0])
    return [item[1] for item in scored]


# ── Delegation ───────────────────────────────────────────────────────────


def delegate_task(
    title: str,
    description: str,
    required_capability: str,
    channel: str = "tasks",
    priority: str = "normal",
) -> dict:
    """Auto-route a task to the best available agent for the required capability.
    Posts the task to the channel and notifies the selected agent."""
    agents = find_agents(capability=required_capability)

    if not agents:
        # No capable agent found — post as unassigned
        result = _post_task(
            title=title,
            description=description,
            channel=channel,
            priority=priority,
            assigned_to=None,
            required_capability=required_capability,
        )
        result["routed"] = False
        result["reason"] = f"No agent with capability '{required_capability}' found"
        return result

    # Pick the best scoring agent
    best = agents[0]
    result = _post_task(
        title=title,
        description=description,
        channel=channel,
        priority=priority,
        assigned_to=best["session_name"],
        assigned_session_id=best["session_id"],
        required_capability=required_capability,
    )
    result["routed"] = True
    result["assigned_to"] = best["session_name"]
    result["assigned_platform"] = best["platform"]
    result["match_score"] = best["score"]

    # Log for load balancing
    _log_delegation(best["session_id"], best["session_name"], title, required_capability)

    return result


def _post_task(
    title: str,
    description: str,
    channel: str,
    priority: str,
    assigned_to: Optional[str],
    assigned_session_id: Optional[str] = None,
    required_capability: Optional[str] = None,
) -> dict:
    """Post a delegated task to a channel."""
    task_id = uuid.uuid4().hex[:12]

    assignment = f" -> assigned to {assigned_to}" if assigned_to else " (unassigned — no capable agent)"
    message = f"[TASK:{task_id}] [{priority.upper()}] {title}{assignment}\n{description}"
    if required_capability:
        message += f"\nRequired capability: {required_capability}"

    entry = {
        "id": task_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "session_id": SESSION_ID,
        "session_name": "brynq-dispatcher",
        "platform": "brynq",
        "channel": channel,
        "msg_type": "request",
        "message": message,
        "task_id": task_id,
        "assigned_to": assigned_to,
        "assigned_session_id": assigned_session_id,
        "required_capability": required_capability,
        "priority": priority,
        "status": "delegated" if assigned_to else "pending",
    }
    _append_jsonl_sync(CHANNELS_DIR / f"{channel}.jsonl", entry)

    # Also create a real task file so broker_accept_task/broker_complete_task work
    from .tasks import TASKS_DIR
    TASKS_DIR.mkdir(parents=True, exist_ok=True)
    task_file = {
        "task_id": task_id,
        "title": title,
        "description": description,
        "channel": channel,
        "priority": priority,
        "status": "delegated" if assigned_to else "pending",
        "requester_session": SESSION_ID,
        "requester_name": SESSION_NAME,
        "assignee_session": assigned_session_id,
        "assignee_name": assigned_to,
        "required_capability": required_capability,
        "created": datetime.now(timezone.utc).isoformat(),
    }
    _write_json_sync(TASKS_DIR / f"{task_id}.json", task_file)

    # Also notify the assigned agent directly via #_system
    if assigned_to:
        notify = {
            "id": uuid.uuid4().hex[:12],
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "session_id": SESSION_ID,
            "session_name": "brynq-dispatcher",
            "platform": "brynq",
            "channel": "_system",
            "msg_type": "request",
            "message": f"@{assigned_to}: You have been assigned task [{task_id}] '{title}' ({priority}). Check #{channel} for details.",
        }
        _append_jsonl_sync(CHANNELS_DIR / "_system.jsonl", notify)

    return {"task_id": task_id, "channel": channel, "priority": priority}


def _log_delegation(session_id: str, session_name: str, title: str, capability: str) -> None:
    """Log a delegation event for load balancing tracking."""
    entry = {
        "session_id": session_id,
        "session_name": session_name,
        "title": title,
        "capability": capability,
        "delegated_at": datetime.now(timezone.utc).isoformat(),
    }
    _append_jsonl_sync(TASK_LOG, entry)


def _get_delegation_counts() -> dict[str, int]:
    """Get active delegation counts per session for load balancing."""
    counts: dict[str, int] = {}
    entries = _read_jsonl_sync(TASK_LOG)
    # Count recent delegations (last 100 entries as proxy)
    for entry in entries[-100:]:
        sid = entry.get("session_id", "")
        counts[sid] = counts.get(sid, 0) + 1
    return counts
