"""
Phase 44: Intelligent Task Routing — Multi-Factor Scoring & Learning.

Builds on Phase 31 (auto_dispatch) and Phase 5 (capability) to add:
  - Multi-factor scoring: skill match, performance, availability, reputation
  - Historical performance tracking per agent-skill pair
  - Workload-aware load balancing
  - Custom routing rules (admin-defined conditions)
  - Fallback chains for high-availability routing
  - Routing analytics and stats

Storage:
  state/broker/routing/
    history.jsonl                     — routing decision log
    affinities/{agent_id}.json       — learned agent-skill performance
    rules.json                       — custom routing rules
"""

import json
import uuid
from datetime import datetime, timezone, timedelta
from pathlib import Path
from typing import Any, Callable, Optional

from brynq.server import (
    BROKER_DIR,
    _write_json_sync,
    _read_jsonl_sync,
    _append_jsonl_sync,
)

# ── Paths ──────────────────────────────────────────────────────────────────

ROUTING_DIR = BROKER_DIR / "routing"
HISTORY_FILE = ROUTING_DIR / "history.jsonl"
AFFINITIES_DIR = ROUTING_DIR / "affinities"
RULES_FILE = ROUTING_DIR / "rules.json"
PROFILES_DIR = BROKER_DIR / "profiles"
TASKS_DIR = BROKER_DIR / "tasks"

for _d in (ROUTING_DIR, AFFINITIES_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# ── In-memory routing rules registry ──────────────────────────────────────

_routing_rules: dict[str, dict] = {}

# ── Fallback chains (in-memory + persisted) ───────────────────────────────

_fallback_chains: dict[str, list[str]] = {}

# ── Scoring weights ───────────────────────────────────────────────────────

WEIGHTS = {
    "skill_match": 0.40,
    "performance": 0.25,
    "availability": 0.20,
    "reputation": 0.15,
}


# ── 1. Route Task ─────────────────────────────────────────────────────────


def route_task(
    task_title: str,
    required_skills: list[str],
    priority: str = "normal",
    constraints: Optional[dict] = None,
) -> list[dict]:
    """Find the best agents for a task using multi-factor scoring.

    Returns a ranked list of candidates, each with:
      agent_id, score, breakdown (skill_match, performance, availability,
      reputation), reasons list.
    """
    constraints = constraints or {}
    profiles = _load_all_profiles()
    if not profiles:
        return []

    candidates = []
    for profile in profiles:
        agent_id = profile.get("session_id", "")

        # Check constraint: excluded agents
        excluded = constraints.get("exclude_agents", [])
        if agent_id in excluded:
            continue

        # Check constraint: required platform
        req_platform = constraints.get("platform")
        if req_platform and profile.get("platform", "") != req_platform:
            continue

        # Skip offline / exiled agents
        avail = profile.get("availability", "available")
        if avail in ("offline", "exiled"):
            continue

        # Evaluate custom routing rules — skip if any rule rejects
        rejected, rule_reasons = _evaluate_rules(profile, task_title, required_skills, priority)
        if rejected:
            continue

        breakdown, reasons = _score_agent(
            profile, required_skills, task_title, priority,
        )

        reasons.extend(rule_reasons)

        total = (
            breakdown["skill_match"] * WEIGHTS["skill_match"]
            + breakdown["performance"] * WEIGHTS["performance"]
            + breakdown["availability"] * WEIGHTS["availability"]
            + breakdown["reputation"] * WEIGHTS["reputation"]
        )

        candidates.append({
            "agent_id": agent_id,
            "agent_name": profile.get("session_name", agent_id),
            "score": round(total, 3),
            "breakdown": {k: round(v, 3) for k, v in breakdown.items()},
            "reasons": reasons,
        })

    # Sort descending by score
    candidates.sort(key=lambda c: c["score"], reverse=True)
    return candidates


def _score_agent(
    profile: dict,
    required_skills: list[str],
    task_title: str,
    priority: str,
) -> tuple[dict, list[str]]:
    """Compute per-factor scores (0-100 each) and collect reasons."""
    agent_id = profile.get("session_id", "")
    reasons: list[str] = []

    # ── Skill match (0-100) ───────────────────────────────────────────
    agent_skills = [s.lower() for s in profile.get("skills", [])]
    req_lower = [s.lower() for s in required_skills]

    if req_lower:
        matched = [s for s in req_lower if s in agent_skills]
        skill_score = (len(matched) / len(req_lower)) * 100
        if matched:
            reasons.append(f"Matches skills: {', '.join(matched)}")
        missing = [s for s in req_lower if s not in agent_skills]
        if missing:
            reasons.append(f"Missing skills: {', '.join(missing)}")
    else:
        skill_score = 50.0  # No skills required = neutral
        reasons.append("No specific skills required")

    # ── Performance / affinity (0-100) ────────────────────────────────
    if req_lower:
        affinities = [get_agent_affinity(agent_id, s) for s in req_lower]
        perf_score = sum(affinities) / len(affinities)
    else:
        perf_score = 50.0  # Neutral when no skills specified

    if perf_score >= 70:
        reasons.append(f"Strong historical performance ({perf_score:.0f})")
    elif perf_score <= 30:
        reasons.append(f"Low historical performance ({perf_score:.0f})")

    # ── Availability / workload (0-100) ───────────────────────────────
    workload = get_workload(agent_id)
    agent_tasks = workload.get(agent_id, 0) if isinstance(workload, dict) else 0
    avail = profile.get("availability", "available")

    if avail == "available":
        avail_score = 100.0
    elif avail == "busy":
        avail_score = 40.0
    else:
        avail_score = 0.0

    # Penalize for each active task (diminishing returns)
    avail_score = max(0.0, avail_score - agent_tasks * 15)

    if agent_tasks == 0:
        reasons.append("No active tasks")
    else:
        reasons.append(f"{agent_tasks} active task(s)")

    # ── Reputation (0-100) ────────────────────────────────────────────
    rep = profile.get("reputation", 3.0)
    rep_score = min(rep / 5.0, 1.0) * 100
    if rep >= 4.5:
        reasons.append(f"Excellent reputation ({rep:.1f})")
    elif rep < 2.5:
        reasons.append(f"Low reputation ({rep:.1f})")

    breakdown = {
        "skill_match": skill_score,
        "performance": perf_score,
        "availability": avail_score,
        "reputation": rep_score,
    }
    return breakdown, reasons


# ── 2. Routing History ─────────────────────────────────────────────────────


def record_routing(
    task_id: str,
    agent_id: str,
    score: float,
    outcome: Optional[str] = None,
) -> dict:
    """Log a routing decision."""
    entry = {
        "id": uuid.uuid4().hex[:12],
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "task_id": task_id,
        "agent_id": agent_id,
        "score": score,
        "outcome": outcome,
    }
    _append_jsonl_sync(HISTORY_FILE, entry)
    return entry


def record_outcome(
    task_id: str,
    success: bool,
    completion_time_minutes: float,
) -> dict:
    """Record how a routing decision turned out.

    Updates the routing history and adjusts the agent's affinity scores.
    """
    entry = {
        "id": uuid.uuid4().hex[:12],
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "task_id": task_id,
        "success": success,
        "completion_time_minutes": completion_time_minutes,
        "type": "outcome",
    }
    _append_jsonl_sync(HISTORY_FILE, entry)

    # Find the original routing entry for this task to update affinities
    history = _read_jsonl_sync(HISTORY_FILE)
    original = None
    for h in history:
        if h.get("task_id") == task_id and h.get("agent_id") and h.get("type") != "outcome":
            original = h
            break

    if original:
        agent_id = original["agent_id"]
        # Try to find the task's skills from the task file
        task_skills = _get_task_skills(task_id)
        for skill in task_skills:
            _update_affinity(agent_id, skill, success, completion_time_minutes)

    return entry


# ── 3. Learn from History — Agent Affinities ──────────────────────────────


def get_agent_affinity(agent_id: str, skill: str) -> float:
    """Return a performance score (0-100) for a specific agent-skill pair.

    Based on historical outcomes. Returns 50.0 (neutral) if no data.
    """
    path = AFFINITIES_DIR / f"{agent_id}.json"
    try:
        data = json.loads(path.read_text("utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return 50.0  # No data = neutral

    skill_data = data.get(skill.lower())
    if not skill_data:
        return 50.0

    total = skill_data.get("successes", 0) + skill_data.get("failures", 0)
    if total == 0:
        return 50.0

    success_rate = skill_data.get("successes", 0) / total
    # Blend with prior: new agents regress toward 50
    confidence = min(total / 10.0, 1.0)  # Full confidence at 10+ tasks
    blended = 50.0 * (1 - confidence) + (success_rate * 100) * confidence

    # Bonus for fast completion
    avg_time = skill_data.get("avg_completion_minutes", 0)
    if avg_time > 0 and avg_time < 30:
        blended = min(100.0, blended + 5)

    return round(blended, 2)


def _update_affinity(
    agent_id: str,
    skill: str,
    success: bool,
    completion_time_minutes: float,
) -> None:
    """Update the affinity data for an agent-skill pair."""
    path = AFFINITIES_DIR / f"{agent_id}.json"
    try:
        data = json.loads(path.read_text("utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        data = {}

    key = skill.lower()
    if key not in data:
        data[key] = {"successes": 0, "failures": 0, "avg_completion_minutes": 0, "total_tasks": 0}

    entry = data[key]
    if success:
        entry["successes"] += 1
    else:
        entry["failures"] += 1

    entry["total_tasks"] = entry["successes"] + entry["failures"]

    # Running average of completion time
    old_total = entry["total_tasks"] - 1
    if old_total > 0:
        old_avg = entry["avg_completion_minutes"]
        entry["avg_completion_minutes"] = round(
            (old_avg * old_total + completion_time_minutes) / entry["total_tasks"], 2
        )
    else:
        entry["avg_completion_minutes"] = round(completion_time_minutes, 2)

    entry["last_updated"] = datetime.now(timezone.utc).isoformat()
    data[key] = entry
    _write_json_sync(path, data)


# ── 4. Load Balancing ──────────────────────────────────────────────────────


def get_workload(agent_id: Optional[str] = None) -> dict[str, int]:
    """Return current active task count per agent.

    If agent_id is given, returns {agent_id: count}.
    Otherwise returns counts for all agents.
    """
    counts: dict[str, int] = {}

    if TASKS_DIR.exists():
        for f in TASKS_DIR.glob("*.json"):
            try:
                task = json.loads(f.read_text("utf-8"))
            except (json.JSONDecodeError, OSError):
                continue

            status = task.get("status", "")
            if status not in ("accepted", "in_progress", "delegated"):
                continue

            assignee = task.get("assignee_session") or task.get("assigned_session_id", "")
            if not assignee:
                continue

            if agent_id and assignee != agent_id:
                continue

            counts[assignee] = counts.get(assignee, 0) + 1

    if agent_id and agent_id not in counts:
        counts[agent_id] = 0

    return counts


def get_least_loaded(skill: Optional[str] = None) -> Optional[dict]:
    """Return the agent with the fewest active tasks.

    Optionally filters to agents that have the specified skill.
    Returns dict with agent_id, agent_name, active_tasks, or None.
    """
    profiles = _load_all_profiles()
    if not profiles:
        return None

    workloads = get_workload()

    candidates = []
    for profile in profiles:
        avail = profile.get("availability", "available")
        if avail in ("offline", "exiled"):
            continue

        agent_id = profile.get("session_id", "")

        if skill:
            agent_skills = [s.lower() for s in profile.get("skills", [])]
            if skill.lower() not in agent_skills:
                continue

        active = workloads.get(agent_id, 0)
        candidates.append({
            "agent_id": agent_id,
            "agent_name": profile.get("session_name", agent_id),
            "active_tasks": active,
        })

    if not candidates:
        return None

    candidates.sort(key=lambda c: c["active_tasks"])
    return candidates[0]


# ── 5. Routing Rules ──────────────────────────────────────────────────────


def add_rule(
    name: str,
    condition_fn: Callable[[dict, str, list[str], str], bool],
    action: str,
) -> dict:
    """Add a custom routing rule.

    Args:
        name: Unique rule name.
        condition_fn: Function(profile, task_title, required_skills, priority)
                      that returns True if the rule applies.
        action: One of 'prefer', 'reject', 'require'.
            - 'prefer': Boost this agent's score.
            - 'reject': Exclude this agent from routing.
            - 'require': Only route to this agent (if condition matches).

    Returns the rule definition.
    """
    rule = {
        "name": name,
        "action": action,
        "created": datetime.now(timezone.utc).isoformat(),
    }
    _routing_rules[name] = {
        "condition_fn": condition_fn,
        "action": action,
        "meta": rule,
    }

    # Persist metadata (not the function — that's in-memory only)
    _persist_rules_meta()

    return rule


def list_rules() -> list[dict]:
    """Return all registered routing rules (metadata only)."""
    return [r["meta"] for r in _routing_rules.values()]


def remove_rule(name: str) -> bool:
    """Remove a routing rule by name. Returns True if found and removed."""
    if name in _routing_rules:
        del _routing_rules[name]
        _persist_rules_meta()
        return True
    return False


def _evaluate_rules(
    profile: dict,
    task_title: str,
    required_skills: list[str],
    priority: str,
) -> tuple[bool, list[str]]:
    """Evaluate all routing rules for a candidate agent.

    Returns (rejected: bool, reasons: list[str]).
    """
    rejected = False
    reasons: list[str] = []

    for name, rule in _routing_rules.items():
        fn = rule["condition_fn"]
        action = rule["action"]
        try:
            matches = fn(profile, task_title, required_skills, priority)
        except Exception:
            continue

        if matches:
            if action == "reject":
                rejected = True
                reasons.append(f"Rejected by rule '{name}'")
            elif action == "prefer":
                reasons.append(f"Preferred by rule '{name}'")
            elif action == "require":
                reasons.append(f"Required by rule '{name}'")

    return rejected, reasons


def _persist_rules_meta() -> None:
    """Persist rule metadata to disk (condition functions are in-memory)."""
    meta = [r["meta"] for r in _routing_rules.values()]
    _write_json_sync(RULES_FILE, {"rules": meta, "updated": datetime.now(timezone.utc).isoformat()})


# ── 6. Routing Stats ──────────────────────────────────────────────────────


def get_routing_stats(hours: int = 24) -> dict:
    """Aggregate routing statistics for the given time window.

    Returns:
        total_routed, success_rate, avg_score, top_agents,
        most_requested_skills, routing_failures.
    """
    cutoff = (datetime.now(timezone.utc) - timedelta(hours=hours)).isoformat()
    history = _read_jsonl_sync(HISTORY_FILE)

    # Filter to time window
    recent = [h for h in history if h.get("timestamp", "") >= cutoff]

    routing_entries = [h for h in recent if h.get("agent_id") and h.get("type") != "outcome"]
    outcome_entries = [h for h in recent if h.get("type") == "outcome"]

    total_routed = len(routing_entries)

    # Success rate from outcomes
    successes = sum(1 for o in outcome_entries if o.get("success"))
    total_outcomes = len(outcome_entries)
    success_rate = round(successes / total_outcomes, 3) if total_outcomes > 0 else 0.0

    # Average score
    scores = [h.get("score", 0) for h in routing_entries if h.get("score") is not None]
    avg_score = round(sum(scores) / len(scores), 3) if scores else 0.0

    # Top agents by routing count
    agent_counts: dict[str, int] = {}
    for h in routing_entries:
        aid = h.get("agent_id", "")
        agent_counts[aid] = agent_counts.get(aid, 0) + 1

    top_agents = sorted(agent_counts.items(), key=lambda x: x[1], reverse=True)[:5]
    top_agents = [{"agent_id": a, "routed_count": c} for a, c in top_agents]

    # Most requested skills — from task files referenced in history
    skill_counts: dict[str, int] = {}
    for h in routing_entries:
        task_skills = _get_task_skills(h.get("task_id", ""))
        for s in task_skills:
            skill_counts[s] = skill_counts.get(s, 0) + 1

    most_requested = sorted(skill_counts.items(), key=lambda x: x[1], reverse=True)[:5]
    most_requested = [{"skill": s, "count": c} for s, c in most_requested]

    # Routing failures = outcomes where success is False
    failures = sum(1 for o in outcome_entries if not o.get("success"))

    return {
        "hours": hours,
        "total_routed": total_routed,
        "success_rate": success_rate,
        "avg_score": avg_score,
        "top_agents": top_agents,
        "most_requested_skills": most_requested,
        "routing_failures": failures,
        "total_outcomes": total_outcomes,
    }


# ── 7. Fallback Chain ─────────────────────────────────────────────────────


def set_fallback_chain(skills: list[str], agent_ids: list[str]) -> dict:
    """Define a fallback order for a set of skills.

    When routing a task requiring these skills, if the top-scored agent is
    unavailable, try agents in the given order.

    Args:
        skills: Skill set this chain covers (sorted and lowered as key).
        agent_ids: Ordered list of agent IDs to try.

    Returns the chain definition.
    """
    key = _chain_key(skills)
    chain = {
        "skills": [s.lower() for s in skills],
        "agent_ids": agent_ids,
        "created": datetime.now(timezone.utc).isoformat(),
    }
    _fallback_chains[key] = agent_ids

    # Persist
    all_chains = _load_fallback_chains()
    all_chains[key] = chain
    _write_json_sync(ROUTING_DIR / "fallback_chains.json", all_chains)

    return chain


def get_fallback_chain(skills: list[str]) -> list[str]:
    """Get the fallback agent order for a skill set.

    Returns agent_ids list, or empty list if no chain defined.
    """
    key = _chain_key(skills)

    # Check in-memory first
    if key in _fallback_chains:
        return _fallback_chains[key]

    # Try disk
    all_chains = _load_fallback_chains()
    chain_data = all_chains.get(key)
    if chain_data:
        return chain_data.get("agent_ids", [])

    return []


def route_with_fallback(
    task_title: str,
    required_skills: list[str],
    priority: str = "normal",
    constraints: Optional[dict] = None,
) -> Optional[dict]:
    """Route a task, falling back through the chain if top agents are busy.

    Returns the selected candidate dict, or None if all fallbacks exhausted.
    """
    candidates = route_task(task_title, required_skills, priority, constraints)
    if not candidates:
        # Try fallback chain
        fallback_ids = get_fallback_chain(required_skills)
        for fid in fallback_ids:
            workload = get_workload(fid)
            if workload.get(fid, 0) < 5:  # Not overloaded
                return {
                    "agent_id": fid,
                    "agent_name": fid,
                    "score": 0.0,
                    "breakdown": {},
                    "reasons": ["Selected via fallback chain"],
                }
        return None

    # Check if top candidate is overloaded; if so, walk down list then fallback
    for candidate in candidates:
        aid = candidate["agent_id"]
        workload = get_workload(aid)
        if workload.get(aid, 0) < 5:
            return candidate

    # All scored candidates overloaded — try fallback chain
    fallback_ids = get_fallback_chain(required_skills)
    for fid in fallback_ids:
        workload = get_workload(fid)
        if workload.get(fid, 0) < 5:
            return {
                "agent_id": fid,
                "agent_name": fid,
                "score": 0.0,
                "breakdown": {},
                "reasons": ["Selected via fallback chain (all scored agents overloaded)"],
            }

    # Return top scored even if overloaded — better than nothing
    return candidates[0] if candidates else None


# ── Internals ──────────────────────────────────────────────────────────────


def _load_all_profiles() -> list[dict]:
    """Load all agent profiles from the profiles directory."""
    profiles = []
    if not PROFILES_DIR.exists():
        return profiles

    for f in PROFILES_DIR.glob("*.json"):
        if f.name.startswith("_"):
            continue
        try:
            profiles.append(json.loads(f.read_text("utf-8")))
        except (json.JSONDecodeError, OSError):
            continue

    return profiles


def _get_task_skills(task_id: str) -> list[str]:
    """Extract skills from a task file, if it exists."""
    if not task_id:
        return []
    path = TASKS_DIR / f"{task_id}.json"
    try:
        task = json.loads(path.read_text("utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return []

    # Check for explicit skills field first
    skills = task.get("required_skills", [])
    if skills:
        return [s.lower() for s in skills]

    # Check required_capability
    cap = task.get("required_capability", "")
    if cap:
        return [cap.lower()]

    return []


def _chain_key(skills: list[str]) -> str:
    """Create a stable key from a skill list."""
    return ",".join(sorted(s.lower() for s in skills))


def _load_fallback_chains() -> dict:
    """Load persisted fallback chains."""
    path = ROUTING_DIR / "fallback_chains.json"
    try:
        return json.loads(path.read_text("utf-8"))
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return {}
