from __future__ import annotations

import asyncio
import json
from enum import Enum
from typing import Any

from pydantic import BaseModel

from nexus.resilience.errors import OrchestratorError


class BranchStrategy(str, Enum):
    FIRST_WINS = "first_wins"
    CONSENSUS = "consensus"
    RANK_AND_SELECT = "rank_and_select"


class BranchResult(BaseModel):
    branch_id: str
    result: Any
    agent_id: str
    duration_ms: float
    tokens_used: int
    quality_score: float | None = None


class BranchManager:
    """Tracks branch forks and resolves them by configured strategy."""

    def __init__(self) -> None:
        self._forks: dict[str, dict[str, Any]] = {}
        self._results: dict[str, list[BranchResult]] = {}

    def register_fork(
        self,
        fork_id: str,
        branch_ids: list[str],
        strategy: BranchStrategy,
        evaluator_fn: Any | None = None,
    ) -> None:
        """Register a fork and expected branch ids."""
        self._forks[fork_id] = {
            "branch_ids": set(branch_ids),
            "strategy": strategy,
            "evaluator_fn": evaluator_fn,
        }
        self._results[fork_id] = []

    def record_result(self, fork_id: str, branch_result: BranchResult) -> None:
        """Record one branch result for a fork."""
        if fork_id not in self._forks:
            raise OrchestratorError(f"Unknown fork id: {fork_id}")
        self._results[fork_id].append(branch_result)

    async def resolve(self, fork_id: str, timeout: float = 300.0) -> tuple[Any, str]:
        """Wait for completion and resolve a fork."""
        if fork_id not in self._forks:
            raise OrchestratorError(f"Unknown fork id: {fork_id}")
        cfg = self._forks[fork_id]

        started = asyncio.get_running_loop().time()
        while not self.is_complete(fork_id):
            if (asyncio.get_running_loop().time() - started) > timeout:
                raise OrchestratorError(f"Timeout resolving fork {fork_id}")
            await asyncio.sleep(0.05)

        results = self._results[fork_id]
        strategy: BranchStrategy = cfg["strategy"]
        if strategy == BranchStrategy.FIRST_WINS:
            first = results[0]
            return first.result, first.branch_id

        if strategy == BranchStrategy.CONSENSUS:
            base = results[0].result
            for r in results[1:]:
                if r.result == base:
                    continue
                if self._string_similarity(base, r.result) < 0.9:
                    raise OrchestratorError(f"Consensus failed for fork {fork_id}")
            return results[0].result, results[0].branch_id

        evaluator_fn = cfg.get("evaluator_fn")
        if evaluator_fn:
            chosen = evaluator_fn(results)
            return chosen.result, chosen.branch_id

        ranked = sorted(results, key=lambda r: (r.quality_score if r.quality_score is not None else -1.0), reverse=True)
        return ranked[0].result, ranked[0].branch_id

    def is_complete(self, fork_id: str) -> bool:
        """Whether all expected branches have produced a result."""
        if fork_id not in self._forks:
            return False
        expected = self._forks[fork_id]["branch_ids"]
        got = {r.branch_id for r in self._results.get(fork_id, [])}
        return expected.issubset(got)

    def get_pending_forks(self) -> list[str]:
        """List unresolved fork ids."""
        return [fork_id for fork_id in self._forks if not self.is_complete(fork_id)]

    def _string_similarity(self, a: Any, b: Any) -> float:
        if a == b:
            return 1.0
        a_s = json.dumps(a, sort_keys=True, default=str)
        b_s = json.dumps(b, sort_keys=True, default=str)
        set_a = set(a_s.split())
        set_b = set(b_s.split())
        if not set_a and not set_b:
            return 1.0
        return len(set_a.intersection(set_b)) / max(1, len(set_a.union(set_b)))
