"""Shared pipeline building blocks for candidate evaluation.

These are **composable utilities**, not a rigid framework.  Each
language package calls them directly, passing language-specific
callables where needed.  Nothing here imports from a language
package.

Typical usage from a language optimizer::

    from codeflash_core._pipeline import (
        CandidateForest,
        EvaluationContext,
        dedup_candidates,
        filter_refined_candidates,
        performance_gain,
        select_best,
    )

    # The language optimizer owns the loop, core provides the logic.
    unique = dedup_candidates(
        candidates,
        normalize_fn=normalize_python_code,
        original_normalized=original_norm,
    )
    ...
    best_id = select_best(eval_ctx, original_runtime_ns, ...)
"""

from __future__ import annotations

import difflib
import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Callable

import attrs

if TYPE_CHECKING:
    from ._model import Candidate

log = logging.getLogger(__name__)


def performance_gain(
    *,
    original_runtime_ns: int,
    optimized_runtime_ns: int,
) -> float:
    """Calculate the speedup of optimized code over the original.

    Returns a ratio where 1.0 means 100% faster (2x speedup).
    Returns 0.0 when the optimized runtime is zero.
    """
    if optimized_runtime_ns == 0:
        return 0.0
    return (original_runtime_ns - optimized_runtime_ns) / optimized_runtime_ns


def diff_length(a: str, b: str) -> int:
    """Compute the character length of a unified diff between two strings."""
    a_lines = a.splitlines(keepends=True)
    b_lines = b.splitlines(keepends=True)
    diff_lines = list(
        difflib.unified_diff(a_lines, b_lines, lineterm=""),
    )
    return len("\n".join(diff_lines))


def create_rank_dictionary(values: list[int]) -> dict[int, int]:
    """Map list indices to their rank in ascending order.

    Returns ``{original_index: rank}`` where rank 0 is smallest.
    """
    sorted_indices = sorted(
        range(len(values)),
        key=lambda i: values[i],
    )
    return {
        original_index: rank
        for rank, original_index in enumerate(sorted_indices)
    }


@attrs.define
class CandidateNode:
    """A node in the refinement tree, wrapping one :class:`Candidate`."""

    candidate: Candidate
    parent: CandidateNode | None = None
    children: list[CandidateNode] = attrs.Factory(list)

    @property
    def candidate_id(self) -> str:
        """Shortcut for the wrapped candidate's id."""
        return self.candidate.candidate_id

    def is_leaf(self) -> bool:
        """Return *True* if this node has no children."""
        return len(self.children) == 0

    def path_to_root(self) -> list[Candidate]:
        """Return the ancestor chain from root down to this node."""
        chain: list[Candidate] = []
        node: CandidateNode | None = self
        while node is not None:
            chain.append(node.candidate)
            node = node.parent
        chain.reverse()
        return chain


@attrs.define
class CandidateForest:
    """A forest of refinement trees tracking parent-child relationships.

    Each tree root is an initial candidate from the AI service.
    Children are refinements, repairs, or adaptive optimizations
    that reference a parent via ``parent_id``.
    """

    _nodes: dict[str, CandidateNode] = attrs.Factory(dict)

    def add(self, candidate: Candidate) -> CandidateNode:
        """Insert *candidate* into the forest, linking to its parent."""
        cid = candidate.candidate_id

        # If this id was used as a placeholder parent, fill it in.
        if cid in self._nodes:
            self._nodes[cid].candidate = candidate
            return self._nodes[cid]

        parent_node: CandidateNode | None = None
        if candidate.parent_id:
            if candidate.parent_id not in self._nodes:
                # Placeholder — the parent hasn't been added yet.
                placeholder = attrs.evolve(
                    candidate,
                    candidate_id=candidate.parent_id,
                )
                self._nodes[candidate.parent_id] = CandidateNode(
                    candidate=placeholder,
                )
            parent_node = self._nodes[candidate.parent_id]

        node = CandidateNode(candidate=candidate, parent=parent_node)
        if parent_node is not None:
            parent_node.children.append(node)
        self._nodes[cid] = node
        return node

    def get(self, candidate_id: str) -> CandidateNode | None:
        """Look up a node by candidate id."""
        return self._nodes.get(candidate_id)

    def __len__(self) -> int:
        """Return the number of nodes in the forest."""
        return len(self._nodes)


@attrs.define
class EvaluationContext:
    """Tracks results across candidate evaluations.

    Each candidate is identified by its ``candidate_id`` string.
    The context accumulates speedup ratios, runtimes, correctness
    flags, and deduplication info so that the selection step can
    pick the best candidate.
    """

    speedup_ratios: dict[str, float | None] = attrs.Factory(dict)
    optimized_runtimes: dict[str, float | None] = attrs.Factory(dict)
    is_correct: dict[str, bool] = attrs.Factory(dict)
    line_profiler_results: dict[str, str] = attrs.Factory(dict)
    code_to_id: dict[str, dict[str, Any]] = attrs.Factory(dict)
    post_comments: dict[str, str] = attrs.Factory(dict)
    valid_candidates: list[Any] = attrs.Factory(list)
    async_throughputs: dict[str, int] = attrs.Factory(dict)
    candidate_concurrency: dict[str, Any] = attrs.Factory(dict)
    optimizations_post: dict[str, str] = attrs.Factory(dict)

    def record_failed(self, candidate_id: str) -> None:
        """Record that a candidate failed behavioral tests."""
        self.optimized_runtimes[candidate_id] = None
        self.is_correct[candidate_id] = False
        self.speedup_ratios[candidate_id] = None

    def record_success(
        self,
        candidate_id: str,
        runtime: float,
        speedup: float,
    ) -> None:
        """Record a passing candidate with its runtime and speedup."""
        self.optimized_runtimes[candidate_id] = runtime
        self.is_correct[candidate_id] = True
        self.speedup_ratios[candidate_id] = speedup

    def record_line_profile(
        self,
        candidate_id: str,
        result: str,
    ) -> None:
        """Store line profiler output for a candidate."""
        self.line_profiler_results[candidate_id] = result

    def register_new(
        self,
        normalized_code: str,
        candidate_id: str,
        flat_code: str,
        original_flat_code: str,
    ) -> None:
        """Register a candidate that hasn't been seen before."""
        self.code_to_id[normalized_code] = {
            "candidate_id": candidate_id,
            "shorter_code": flat_code,
            "diff_len": diff_length(flat_code, original_flat_code),
        }

    def handle_duplicate(
        self,
        candidate_id: str,
        normalized_code: str,
        original_flat_code: str,
        flat_code: str,
    ) -> None:
        """Copy prior results for a duplicate candidate.

        If the new candidate has a shorter diff, update the
        stored ``shorter_code``.
        """
        entry = self.code_to_id[normalized_code]
        prior_id = entry["candidate_id"]
        self.speedup_ratios[candidate_id] = self.speedup_ratios.get(
            prior_id,
        )
        self.is_correct[candidate_id] = self.is_correct.get(
            prior_id,
            False,
        )
        self.optimized_runtimes[candidate_id] = self.optimized_runtimes.get(
            prior_id
        )
        if prior_id in self.line_profiler_results:
            self.line_profiler_results[candidate_id] = (
                self.line_profiler_results[prior_id]
            )
        new_diff = diff_length(flat_code, original_flat_code)
        if new_diff < entry["diff_len"]:
            entry["shorter_code"] = flat_code
            entry["diff_len"] = new_diff

    def get_speedup(self, candidate_id: str) -> float | None:
        """Return the speedup ratio, or *None* if not recorded."""
        return self.speedup_ratios.get(candidate_id)

    def get_runtime(self, candidate_id: str) -> float | None:
        """Return the optimized runtime, or *None* if not recorded."""
        return self.optimized_runtimes.get(candidate_id)


def dedup_candidates(
    candidates: list[Candidate],
    *,
    normalize_fn: Callable[[str], str],
    original_normalized: str,
    seen: set[str] | None = None,
    cross_batch: dict[str, dict[str, Any]] | None = None,
) -> list[Candidate]:
    """Remove duplicate candidates using code normalization.

    Candidates are dropped if they:

    * normalize to the same code as the original,
    * were already seen in a prior batch (*cross_batch*), or
    * duplicate another candidate in this batch (*seen*).

    *normalize_fn* is language-specific (e.g.
    :func:`normalize_python_code`).

    Returns the unique candidates in their original order.
    """
    if seen is None:
        seen = set()
    if cross_batch is None:
        cross_batch = {}

    unique: list[Candidate] = []
    removed_original = 0
    removed_cross = 0
    removed_dup = 0

    for candidate in candidates:
        try:
            normalized = normalize_fn(candidate.code)
        except Exception:  # noqa: BLE001
            log.debug(
                "Failed to normalize candidate %s, keeping it",
                candidate.candidate_id,
            )
            unique.append(candidate)
            continue

        if normalized == original_normalized:
            removed_original += 1
            continue
        if normalized in cross_batch:
            removed_cross += 1
            continue
        if normalized in seen:
            removed_dup += 1
            continue

        seen.add(normalized)
        unique.append(candidate)

    if removed_original or removed_cross or removed_dup:
        log.debug(
            "Dedup: removed %d identical-to-original, "
            "%d cross-batch duplicates, %d intra-batch duplicates "
            "from %d candidates",
            removed_original,
            removed_cross,
            removed_dup,
            len(candidates),
        )

    return unique


def filter_refined_candidates(
    candidates: list[Candidate],
    eval_ctx: EvaluationContext,
    forest: CandidateForest,
    original_flat_code: str,
    *,
    max_candidates: int = 5,
) -> list[Candidate]:
    """Rank refined candidates and keep the top *max_candidates*.

    Ranking uses a weighted combination of parent runtime (lower
    is better) and diff length (shorter is better).  Candidates
    without a parent runtime are sorted to the end.
    """
    if len(candidates) <= max_candidates:
        return candidates

    scored: list[tuple[float, int, Candidate]] = []
    for candidate in candidates:
        parent_runtime: float | None = None
        if candidate.parent_id:
            parent_runtime = eval_ctx.get_runtime(candidate.parent_id)

        diff_len = diff_length(candidate.code, original_flat_code)

        # Candidates with unknown parent runtime get worst score.
        runtime_val = parent_runtime if parent_runtime is not None else 1e18
        scored.append((runtime_val, diff_len, candidate))

    if not scored:
        return candidates

    # Normalize both dimensions to [0, 1].
    runtimes = [s[0] for s in scored]
    diffs = [s[1] for s in scored]
    rt_min, rt_max = min(runtimes), max(runtimes)
    df_min, df_max = min(diffs), max(diffs)
    rt_range = rt_max - rt_min or 1.0
    df_range = df_max - df_min or 1.0

    runtime_weight = 0.6
    diff_weight = 0.4

    combined: list[tuple[float, int, Candidate]] = []
    for i, (rt, df, cand) in enumerate(scored):
        norm_rt = (rt - rt_min) / rt_range
        norm_df = (df - df_min) / df_range
        score = runtime_weight * norm_rt + diff_weight * norm_df
        combined.append((score, i, cand))

    combined.sort(key=lambda x: (x[0], x[1]))
    return [cand for _, _, cand in combined[:max_candidates]]


_RUNTIME_WEIGHT = 2
_DIFF_WEIGHT = 1


def select_best(
    eval_ctx: EvaluationContext,
    original_runtime_ns: int,
    diff_lengths: list[int],
    candidate_ids: list[str],
) -> str | None:
    """Pick the best candidate by weighted rank-sum.

    Ranks candidates by runtime (faster = better) and diff length
    (shorter = better), then combines with ``_RUNTIME_WEIGHT`` and
    ``_DIFF_WEIGHT``.  Runtime is weighted more heavily because
    speed is the primary optimization goal.

    Returns the ``candidate_id`` of the winner, or *None*.
    """
    if not candidate_ids:
        return None
    if len(candidate_ids) == 1:
        return candidate_ids[0]

    runtimes: list[int] = []
    for cid in candidate_ids:
        runtime = eval_ctx.get_runtime(cid)
        runtimes.append(
            int(runtime) if runtime is not None else original_runtime_ns,
        )

    diff_ranking = create_rank_dictionary(diff_lengths)
    runtime_ranking = create_rank_dictionary(runtimes)
    overall = {
        k: _DIFF_WEIGHT * diff_ranking[k]
        + _RUNTIME_WEIGHT * runtime_ranking[k]
        for k in diff_ranking
    }
    best_idx = min(overall, key=overall.get)  # type: ignore[arg-type]
    return candidate_ids[best_idx]
