"""Candidate ranking, AST-level dedup, and best-candidate selection."""

from __future__ import annotations

import ast
import difflib
import logging
from functools import lru_cache
from typing import Any, TypeVar

import attrs

log = logging.getLogger(__name__)

ASTNodeT = TypeVar("ASTNodeT", bound=ast.AST)


def normalize_node(node: ASTNodeT) -> ASTNodeT:
    """Strip docstrings and imports from an AST node recursively."""
    if isinstance(
        node, (ast.Module, ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
    ) and ast.get_docstring(node):
        node.body = node.body[1:]
    if hasattr(node, "body"):
        node.body = [
            normalize_node(n)
            for n in node.body
            if not isinstance(n, (ast.Import, ast.ImportFrom))
        ]
    return node


@lru_cache(maxsize=3)
def normalize_code(code: str) -> str:
    """Parse, normalize, and unparse code for AST-level dedup."""
    try:
        return ast.unparse(normalize_node(ast.parse(code)))
    except SyntaxError:
        log.warning("Failed to normalize candidate code, using raw string")
        return code


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


def create_rank_dictionary_compact(
    int_array: list[int],
) -> dict[int, int]:
    """Map original indices to their ranks in ascending order."""
    sorted_indices = sorted(range(len(int_array)), key=lambda i: int_array[i])
    return {
        original_index: rank
        for rank, original_index in enumerate(sorted_indices)
    }


@attrs.define
class CandidateEvaluationContext:
    """Holds tracking state during candidate evaluation."""

    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)
    optimized_line_profiler_results: dict[str, str] = attrs.Factory(dict)
    ast_code_to_id: dict[str, dict[str, Any]] = attrs.Factory(dict)
    optimizations_post: dict[str, str] = attrs.Factory(dict)
    valid_optimizations: list[Any] = attrs.Factory(list)

    def record_failed_candidate(self, optimization_id: str) -> None:
        """Record results for a failed candidate."""
        self.optimized_runtimes[optimization_id] = None
        self.is_correct[optimization_id] = False
        self.speedup_ratios[optimization_id] = None

    def record_successful_candidate(
        self,
        optimization_id: str,
        runtime: float,
        speedup: float,
    ) -> None:
        """Record results for a successful candidate."""
        self.optimized_runtimes[optimization_id] = runtime
        self.is_correct[optimization_id] = True
        self.speedup_ratios[optimization_id] = speedup

    def record_line_profiler_result(
        self,
        optimization_id: str,
        result: str,
    ) -> None:
        """Record line profiler results for a candidate."""
        self.optimized_line_profiler_results[optimization_id] = result

    def handle_duplicate_candidate(
        self,
        optimization_id: str,
        normalized_code: str,
        original_flat_code: str,
        candidate_source_code_flat: str,
    ) -> None:
        """Handle a duplicate candidate by copying prior results."""
        past_opt_id = self.ast_code_to_id[normalized_code]["optimization_id"]
        self.speedup_ratios[optimization_id] = self.speedup_ratios.get(
            past_opt_id
        )
        self.is_correct[optimization_id] = self.is_correct.get(
            past_opt_id, False
        )
        self.optimized_runtimes[optimization_id] = self.optimized_runtimes.get(
            past_opt_id
        )
        if past_opt_id in self.optimized_line_profiler_results:
            self.optimized_line_profiler_results[optimization_id] = (
                self.optimized_line_profiler_results[past_opt_id]
            )
        # Update to shorter code if this candidate has a shorter diff
        new_diff_len = diff_length(
            candidate_source_code_flat, original_flat_code
        )
        if new_diff_len < self.ast_code_to_id[normalized_code]["diff_len"]:
            self.ast_code_to_id[normalized_code]["shorter_source_code"] = (
                candidate_source_code_flat
            )
            self.ast_code_to_id[normalized_code]["diff_len"] = new_diff_len

    def register_new_candidate(
        self,
        normalized_code: str,
        optimization_id: str,
        source_code_flat: str,
        original_flat_code: str,
    ) -> None:
        """Register a new candidate that hasn't been seen before."""
        self.ast_code_to_id[normalized_code] = {
            "optimization_id": optimization_id,
            "shorter_source_code": source_code_flat,
            "diff_len": diff_length(source_code_flat, original_flat_code),
        }

    def get_speedup_ratio(self, optimization_id: str) -> float | None:
        """Return the speedup ratio for the given optimization, or None."""
        return self.speedup_ratios.get(optimization_id)

    def get_optimized_runtime(self, optimization_id: str) -> float | None:
        """Return the optimized runtime for the given optimization, or None."""
        return self.optimized_runtimes.get(optimization_id)


def select_best_candidate(
    eval_ctx: CandidateEvaluationContext,
    original_runtime_ns: int,
    diff_lengths: list[int],
    optimization_ids: list[str],
) -> str | None:
    """Select the best candidate by combined diff and runtime rank.

    Returns the optimization_id of the best candidate,
    or None if no valid candidates.
    """
    if not optimization_ids:
        return None
    if len(optimization_ids) == 1:
        return optimization_ids[0]

    # Build runtime list from eval_ctx
    runtimes = []
    for opt_id in optimization_ids:
        runtime = eval_ctx.get_optimized_runtime(opt_id)
        runtimes.append(
            int(runtime) if runtime is not None else original_runtime_ns
        )

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