"""Test-Time Compute pour du code : best-of-N, arbre de réparations (beam) et MCTS, avec pytest comme juge exact.

Dans un MCTS classique, la valeur d'une feuille s'estime par des simulations. Ici, chaque nœud est un programme
complet que pytest note exactement (part des tests réussis) : pas de simulation, l'évaluation EST le jugement.
"""
from __future__ import annotations

import hashlib
import json
import math
from dataclasses import dataclass, field
from typing import Optional

from mouh.verification.generators import parse_files

from ..contracts import Candidate


def signature(files: dict) -> str:
    return hashlib.sha256(json.dumps(files, sort_keys=True).encode()).hexdigest()


def _files(text: str) -> dict:
    return {k: v for k, v in parse_files(text).items() if not k.startswith("test_")}


class _Search:
    """Comptabilité commune : budget de générations, déduplication, notation."""

    def __init__(self, backend, verifier, tests: dict, budget: int) -> None:
        self.backend, self.verifier, self.tests, self.budget = backend, verifier, tests, budget
        self.spent = 0
        self.seen: dict[str, Candidate] = {}

    @property
    def done(self) -> bool:
        return self.spent >= self.budget or any(c.ok for c in self.seen.values())

    def expand(self, messages, n: int, depth: int, origin: str, greedy_first: bool = True) -> list[Candidate]:
        n = max(0, min(n, self.budget - self.spent))
        if n == 0:
            return []
        out = []
        for text in self.backend.complete(messages, n, greedy_first=greedy_first):
            self.spent += 1
            files = _files(text)
            sig = signature(files) if files else ""
            if files and sig not in self.seen:
                c = self.verifier.score(files, self.tests)
                c.origin, c.depth = origin, depth
                self.seen[sig] = c
                out.append(c)
        return out

    def ranked(self) -> list[Candidate]:
        return sorted(self.seen.values(), key=Candidate.key, reverse=True)


class BestOfN:
    """Une génération gloutonne + des échantillons ; la référence à battre, à budget égal."""
    name = "best_of_n"

    def solve(self, request, plan, backend, verifier, build) -> list[Candidate]:
        tests = dict(request.task.test_files)
        if not tests:  # rien pour juger : une seule proposition, comme la v3.8
            files = _files(backend.complete(build(), 1)[0])
            return [Candidate(files, "racine")] if files else []
        s = _Search(backend, verifier, tests, plan.budget)
        s.expand(build(), plan.budget, 0, "racine")
        return s.ranked()


class RepairBeam:
    """Arbre de réparations en faisceau : à chaque niveau, les `width` meilleurs sont réparés avec leur trace."""
    name = "beam"

    def __init__(self, width: int = 2, depth: int = 3) -> None:
        self.width, self.depth = width, depth

    def solve(self, request, plan, backend, verifier, build) -> list[Candidate]:
        tests = dict(request.task.test_files)
        if not tests:
            return BestOfN().solve(request, plan, backend, verifier, build)
        s = _Search(backend, verifier, tests, plan.budget)
        frontier = s.expand(build(), max(1, plan.budget // 2), 0, "racine")
        for d in range(1, self.depth + 1):
            if s.done or not frontier:
                break
            parents = sorted(frontier, key=Candidate.key, reverse=True)[: self.width]
            frontier = []
            for p in parents:
                frontier += s.expand(build(previous=p.files, diagnostics=[p.output]), 1, d, "reparation")
        return s.ranked()


@dataclass
class _Node:
    cand: Optional[Candidate]                  # None pour la racine (l'énoncé seul)
    parent: Optional["_Node"] = None
    children: list = field(default_factory=list)
    visits: int = 0
    total: float = 0.0
    expansions: int = 0                        # générations déjà tirées depuis ce nœud

    @property
    def q(self) -> float:
        return self.total / self.visits if self.visits else 0.0


class RepairMCTS:
    """MCTS (sélection UCT) sur l'arbre des réparations.

    - Sélection : parmi la racine (nouvelle proposition depuis l'énoncé) et tous les candidats déjà notés,
      on choisit celui qui maximise Q + c * sqrt(ln N / (1 + n)).
    - Expansion : la racine produit un nouvel échantillon ; un candidat est réparé avec sa propre sortie pytest
      (1re réparation gloutonne, les suivantes échantillonnées).
    - Évaluation : pytest, exact (part des tests réussis).
    - Rétropropagation : la valeur remonte jusqu'à la racine.
    Arrêt au premier candidat qui passe tous les tests ou quand le budget est dépensé.
    """
    name = "mcts"

    def __init__(self, c: float = 1.0, root_samples: int = 2, max_children: int = 3) -> None:
        self.c, self.root_samples, self.max_children = c, root_samples, max_children

    def solve(self, request, plan, backend, verifier, build) -> list[Candidate]:
        tests = dict(request.task.test_files)
        if not tests:
            return BestOfN().solve(request, plan, backend, verifier, build)
        s = _Search(backend, verifier, tests, plan.budget)
        root = _Node(None)
        nodes: list[_Node] = []

        def attach(parent: _Node, cands: list[Candidate]) -> None:
            for c in cands:
                node = _Node(c, parent)
                parent.children.append(node)
                nodes.append(node)
                v, cur = c.value, node
                while cur is not None:
                    cur.visits += 1
                    cur.total += v
                    cur = cur.parent

        first = s.expand(build(), min(self.root_samples, plan.budget), 0, "racine")
        root.expansions += min(self.root_samples, plan.budget)
        attach(root, first)
        while not s.done:
            pool = [root] + [n for n in nodes if n.expansions < self.max_children and not n.cand.ok]
            total = max(1, root.visits)

            def uct(n: _Node) -> float:  # Q = moyenne du sous-arbre (le nœud compris)
                return n.q + self.c * math.sqrt(math.log(total + 1) / (1 + n.expansions))

            target = max(pool, key=uct)
            if target is root:
                found = s.expand(build(), 1, 0, "racine", greedy_first=False)
            else:
                found = s.expand(build(previous=target.cand.files, diagnostics=[target.cand.output]), 1,
                                 target.cand.depth + 1, "reparation", greedy_first=target.expansions == 0)
            target.expansions += 1
            if found:
                attach(target, found)
            else:  # doublon ou sortie illisible : la visite compte quand même, pour ne pas reboucler
                cur = target
                while cur is not None:
                    cur.visits += 1
                    cur = cur.parent
        return s.ranked()


class AdaptiveStrategy:
    """Choisit la stratégie selon le plan du routeur : le raisonnement hybride."""
    name = "adaptatif"

    def __init__(self) -> None:
        self.by_mode = {"direct": BestOfN(), "arbre": RepairBeam(), "mcts": RepairMCTS()}

    def solve(self, request, plan, backend, verifier, build) -> list[Candidate]:
        return self.by_mode.get(plan.mode, self.by_mode["direct"]).solve(request, plan, backend, verifier, build)