from __future__ import annotations

import attrs
import pytest

from codeflash_core._model import Candidate
from codeflash_core._pipeline import (
    CandidateForest,
    CandidateNode,
    EvaluationContext,
    create_rank_dictionary,
    dedup_candidates,
    diff_length,
    filter_refined_candidates,
    performance_gain,
    select_best,
)


@pytest.fixture(name="candidate")
def _candidate():
    """
    A minimal Candidate for use in tests.
    """
    return Candidate(
        code="def f(): pass",
        explanation="no-op",
        candidate_id="c1",
    )


@pytest.fixture(name="eval_ctx")
def _eval_ctx():
    """
    A fresh EvaluationContext.
    """
    return EvaluationContext()


class TestPerformanceGain:
    """Tests for performance_gain."""

    def test_two_x_speedup(self):
        """
        Original twice as slow yields a gain of 1.0.
        """
        assert 1.0 == performance_gain(
            original_runtime_ns=200, optimized_runtime_ns=100
        )

    def test_no_speedup(self):
        """
        Equal runtimes yield a gain of 0.0.
        """
        assert 0.0 == performance_gain(
            original_runtime_ns=100, optimized_runtime_ns=100
        )

    def test_slowdown(self):
        """
        Optimized slower than original yields a negative gain.
        """
        assert -0.5 == performance_gain(
            original_runtime_ns=100, optimized_runtime_ns=200
        )

    def test_optimized_zero_returns_zero(self):
        """
        Zero optimized runtime returns 0.0 to avoid division by zero.
        """
        assert 0.0 == performance_gain(
            original_runtime_ns=100, optimized_runtime_ns=0
        )

    def test_both_zero(self):
        """
        Both runtimes zero returns 0.0.
        """
        assert 0.0 == performance_gain(
            original_runtime_ns=0, optimized_runtime_ns=0
        )

    def test_original_zero_optimized_nonzero(self):
        """
        Zero original with nonzero optimized yields -1.0.
        """
        assert -1.0 == performance_gain(
            original_runtime_ns=0, optimized_runtime_ns=100
        )

    def test_large_speedup(self):
        """
        Large runtime difference produces correspondingly large gain.
        """
        assert 999.0 == performance_gain(
            original_runtime_ns=1_000_000, optimized_runtime_ns=1_000
        )


class TestDiffLength:
    """Tests for diff_length."""

    def test_identical_strings(self):
        """
        Identical strings produce a zero-length diff.
        """
        assert 0 == diff_length("hello\n", "hello\n")

    def test_empty_strings(self):
        """
        Two empty strings produce a zero-length diff.
        """
        assert 0 == diff_length("", "")

    def test_one_line_changed(self):
        """
        A single changed line produces a non-zero diff.
        """
        result = diff_length("line1\nline2\n", "line1\nchanged\n")
        assert result > 0

    def test_addition(self):
        """
        Adding a line produces a non-zero diff.
        """
        result = diff_length("a\n", "a\nb\n")
        assert result > 0

    def test_deletion(self):
        """
        Removing a line produces a non-zero diff.
        """
        result = diff_length("a\nb\n", "a\n")
        assert result > 0

    def test_completely_different(self):
        """
        Completely different strings produce a larger diff than a small edit.
        """
        small = diff_length("aaa\n", "aab\n")
        large = diff_length("aaa\nbbb\nccc\n", "xxx\nyyy\nzzz\n")
        assert large > small


class TestCreateRankDictionary:
    """Tests for create_rank_dictionary."""

    def test_ascending_order(self):
        """
        Already-sorted values get ranks 0, 1, 2.
        """
        assert {0: 0, 1: 1, 2: 2} == create_rank_dictionary([10, 20, 30])

    def test_descending_order(self):
        """
        Reverse-sorted values get inverted ranks.
        """
        assert {0: 2, 1: 1, 2: 0} == create_rank_dictionary([30, 20, 10])

    def test_single_element(self):
        """
        A single element gets rank 0.
        """
        assert {0: 0} == create_rank_dictionary([42])

    def test_empty_list(self):
        """
        An empty list returns an empty dict.
        """
        assert {} == create_rank_dictionary([])

    def test_duplicate_values(self):
        """
        Duplicates still get distinct ranks based on original index.
        """
        result = create_rank_dictionary([5, 5, 5])
        assert {0, 1, 2} == set(result.values())


class TestCandidateNode:
    """Tests for CandidateNode."""

    def test_candidate_id_delegates(self, candidate):
        """
        candidate_id property returns the wrapped candidate's id.
        """
        node = CandidateNode(candidate=candidate)
        assert "c1" == node.candidate_id

    def test_is_leaf_true(self, candidate):
        """
        A node with no children is a leaf.
        """
        node = CandidateNode(candidate=candidate)
        assert node.is_leaf() is True

    def test_is_leaf_false(self, candidate):
        """
        A node with children is not a leaf.
        """
        parent = CandidateNode(candidate=candidate)
        child_cand = Candidate(
            code="def g(): pass",
            explanation="child",
            candidate_id="c2",
        )
        child = CandidateNode(candidate=child_cand, parent=parent)
        parent.children.append(child)
        assert parent.is_leaf() is False

    def test_path_to_root_single(self, candidate):
        """
        A root node's path is just itself.
        """
        node = CandidateNode(candidate=candidate)
        path = node.path_to_root()
        assert 1 == len(path)
        assert candidate is path[0]

    def test_path_to_root_chain(self):
        """
        A three-deep chain returns root-first order.
        """
        c1 = Candidate(code="a", explanation="", candidate_id="c1")
        c2 = Candidate(code="b", explanation="", candidate_id="c2")
        c3 = Candidate(code="c", explanation="", candidate_id="c3")
        root = CandidateNode(candidate=c1)
        mid = CandidateNode(candidate=c2, parent=root)
        leaf = CandidateNode(candidate=c3, parent=mid)

        path = leaf.path_to_root()
        assert ["c1", "c2", "c3"] == [c.candidate_id for c in path]

    def test_default_children_empty(self, candidate):
        """
        A new node starts with an empty children list.
        """
        node = CandidateNode(candidate=candidate)
        assert [] == node.children


class TestCandidateForest:
    """Tests for CandidateForest."""

    def test_add_and_get(self, candidate):
        """
        Adding a candidate allows retrieval by id.
        """
        forest = CandidateForest()
        forest.add(candidate)
        node = forest.get("c1")
        assert node is not None
        assert "c1" == node.candidate_id

    def test_len(self):
        """
        __len__ returns the number of nodes in the forest.
        """
        forest = CandidateForest()
        assert 0 == len(forest)
        forest.add(Candidate(code="a", explanation="", candidate_id="c1"))
        assert 1 == len(forest)

    def test_get_missing(self):
        """
        Getting a nonexistent id returns None.
        """
        forest = CandidateForest()
        assert forest.get("nope") is None

    def test_parent_child_linking(self):
        """
        A candidate with parent_id links to its parent node.
        """
        forest = CandidateForest()
        parent = Candidate(code="a", explanation="", candidate_id="p1")
        child = Candidate(
            code="b",
            explanation="",
            candidate_id="c1",
            parent_id="p1",
        )
        forest.add(parent)
        forest.add(child)

        child_node = forest.get("c1")
        assert child_node is not None
        assert child_node.parent is not None
        assert "p1" == child_node.parent.candidate_id

    def test_child_added_before_parent(self):
        """
        Adding a child before its parent creates a placeholder.
        """
        forest = CandidateForest()
        child = Candidate(
            code="b",
            explanation="",
            candidate_id="c1",
            parent_id="p1",
        )
        forest.add(child)
        assert 2 == len(forest)

        parent = Candidate(code="a", explanation="", candidate_id="p1")
        forest.add(parent)
        assert 2 == len(forest)

        parent_node = forest.get("p1")
        assert parent_node is not None
        assert "a" == parent_node.candidate.code

    def test_multiple_roots(self):
        """
        Candidates without parent_id are independent roots.
        """
        forest = CandidateForest()
        forest.add(Candidate(code="a", explanation="", candidate_id="r1"))
        forest.add(Candidate(code="b", explanation="", candidate_id="r2"))
        assert 2 == len(forest)
        assert forest.get("r1").parent is None
        assert forest.get("r2").parent is None


class TestEvaluationContext:
    """Tests for EvaluationContext."""

    def test_record_failed(self, eval_ctx):
        """
        record_failed sets correct flags and None values.
        """
        eval_ctx.record_failed("c1")
        assert eval_ctx.is_correct["c1"] is False
        assert eval_ctx.optimized_runtimes["c1"] is None
        assert eval_ctx.speedup_ratios["c1"] is None

    def test_record_success(self, eval_ctx):
        """
        record_success stores runtime, speedup, and correctness.
        """
        eval_ctx.record_success("c1", runtime=500.0, speedup=1.5)
        assert eval_ctx.is_correct["c1"] is True
        assert 500.0 == eval_ctx.optimized_runtimes["c1"]
        assert 1.5 == eval_ctx.speedup_ratios["c1"]

    def test_get_speedup_missing(self, eval_ctx):
        """
        get_speedup returns None for unknown candidates.
        """
        assert eval_ctx.get_speedup("unknown") is None

    def test_get_runtime_missing(self, eval_ctx):
        """
        get_runtime returns None for unknown candidates.
        """
        assert eval_ctx.get_runtime("unknown") is None

    def test_get_speedup_after_record(self, eval_ctx):
        """
        get_speedup returns the recorded value.
        """
        eval_ctx.record_success("c1", runtime=100.0, speedup=2.0)
        assert 2.0 == eval_ctx.get_speedup("c1")

    def test_get_runtime_after_record(self, eval_ctx):
        """
        get_runtime returns the recorded value.
        """
        eval_ctx.record_success("c1", runtime=100.0, speedup=2.0)
        assert 100.0 == eval_ctx.get_runtime("c1")

    def test_record_line_profile(self, eval_ctx):
        """
        record_line_profile stores the result string.
        """
        eval_ctx.record_line_profile("c1", "profile output")
        assert "profile output" == eval_ctx.line_profiler_results["c1"]

    def test_register_new(self, eval_ctx):
        """
        register_new stores code-to-id mapping with diff info.
        """
        eval_ctx.register_new(
            normalized_code="norm",
            candidate_id="c1",
            flat_code="def f(): pass",
            original_flat_code="def f(): return 1",
        )
        entry = eval_ctx.code_to_id["norm"]
        assert "c1" == entry["candidate_id"]
        assert "def f(): pass" == entry["shorter_code"]
        assert entry["diff_len"] > 0

    def test_handle_duplicate_copies_prior_results(self, eval_ctx):
        """
        handle_duplicate propagates prior correctness and runtime.
        """
        eval_ctx.record_success("c1", runtime=100.0, speedup=2.0)
        eval_ctx.register_new(
            normalized_code="norm",
            candidate_id="c1",
            flat_code="def f(): pass",
            original_flat_code="def g(): pass",
        )

        eval_ctx.handle_duplicate(
            candidate_id="c2",
            normalized_code="norm",
            original_flat_code="def g(): pass",
            flat_code="def f(): pass",
        )

        assert eval_ctx.is_correct["c2"] is True
        assert 100.0 == eval_ctx.optimized_runtimes["c2"]
        assert 2.0 == eval_ctx.speedup_ratios["c2"]

    def test_handle_duplicate_updates_shorter_code(self, eval_ctx):
        """
        handle_duplicate replaces shorter_code when the new diff is smaller.
        """
        eval_ctx.register_new(
            normalized_code="norm",
            candidate_id="c1",
            flat_code="def f():\n    x = 1\n    return x\n",
            original_flat_code="def f():\n    return 1\n",
        )
        eval_ctx.record_success("c1", runtime=100.0, speedup=2.0)

        eval_ctx.handle_duplicate(
            candidate_id="c2",
            normalized_code="norm",
            original_flat_code="def f():\n    return 1\n",
            flat_code="def f():\n    return 1\n",
        )

        assert (
            "def f():\n    return 1\n"
            == eval_ctx.code_to_id["norm"]["shorter_code"]
        )

    def test_starts_empty(self, eval_ctx):
        """
        A new EvaluationContext has empty containers.
        """
        assert {} == eval_ctx.speedup_ratios
        assert {} == eval_ctx.optimized_runtimes
        assert {} == eval_ctx.is_correct
        assert [] == eval_ctx.valid_candidates


class TestDedupCandidates:
    """Tests for dedup_candidates."""

    def test_empty_input(self):
        """
        An empty list returns an empty list.
        """
        result = dedup_candidates(
            [],
            normalize_fn=lambda c: c,
            original_normalized="original",
        )
        assert [] == result

    def test_removes_identical_to_original(self):
        """
        Candidates matching the original are removed.
        """
        c1 = Candidate(code="original", explanation="", candidate_id="c1")
        result = dedup_candidates(
            [c1],
            normalize_fn=lambda c: c,
            original_normalized="original",
        )
        assert [] == result

    def test_removes_intra_batch_duplicates(self):
        """
        Only the first occurrence of a normalized form is kept.
        """
        c1 = Candidate(code="opt", explanation="", candidate_id="c1")
        c2 = Candidate(code="opt", explanation="", candidate_id="c2")
        result = dedup_candidates(
            [c1, c2],
            normalize_fn=lambda c: c,
            original_normalized="original",
        )
        assert 1 == len(result)
        assert "c1" == result[0].candidate_id

    def test_removes_cross_batch_duplicates(self):
        """
        Candidates already in cross_batch are removed.
        """
        c1 = Candidate(code="opt", explanation="", candidate_id="c1")
        result = dedup_candidates(
            [c1],
            normalize_fn=lambda c: c,
            original_normalized="original",
            cross_batch={"opt": {"candidate_id": "prior"}},
        )
        assert [] == result

    def test_keeps_unique_candidates(self):
        """
        All unique candidates are preserved in order.
        """
        c1 = Candidate(code="a", explanation="", candidate_id="c1")
        c2 = Candidate(code="b", explanation="", candidate_id="c2")
        result = dedup_candidates(
            [c1, c2],
            normalize_fn=lambda c: c,
            original_normalized="original",
        )
        assert 2 == len(result)
        assert ["c1", "c2"] == [c.candidate_id for c in result]

    def test_normalize_fn_failure_keeps_candidate(self):
        """
        A candidate whose normalization raises is still kept.
        """

        def bad_normalize(code):
            raise ValueError("boom")

        c1 = Candidate(code="a", explanation="", candidate_id="c1")
        result = dedup_candidates(
            [c1],
            normalize_fn=bad_normalize,
            original_normalized="original",
        )
        assert 1 == len(result)
        assert "c1" == result[0].candidate_id

    def test_seen_set_is_populated(self):
        """
        The seen set accumulates normalized codes across calls.
        """
        seen: set[str] = set()
        c1 = Candidate(code="a", explanation="", candidate_id="c1")
        dedup_candidates(
            [c1],
            normalize_fn=lambda c: c,
            original_normalized="original",
            seen=seen,
        )
        assert "a" in seen

    def test_seen_set_prevents_second_batch(self):
        """
        A code seen in a prior call is treated as a duplicate.
        """
        seen: set[str] = {"a"}
        c1 = Candidate(code="a", explanation="", candidate_id="c1")
        result = dedup_candidates(
            [c1],
            normalize_fn=lambda c: c,
            original_normalized="original",
            seen=seen,
        )
        assert [] == result


class TestFilterRefinedCandidates:
    """Tests for filter_refined_candidates."""

    def test_under_limit_returns_all(self, eval_ctx):
        """
        Fewer candidates than max_candidates returns all of them.
        """
        forest = CandidateForest()
        candidates = [
            Candidate(code="a", explanation="", candidate_id="c1"),
            Candidate(code="b", explanation="", candidate_id="c2"),
        ]
        result = filter_refined_candidates(
            candidates, eval_ctx, forest, "original", max_candidates=5
        )
        assert 2 == len(result)

    def test_at_limit_returns_all(self, eval_ctx):
        """
        Exactly max_candidates returns all of them.
        """
        forest = CandidateForest()
        candidates = [
            Candidate(code="a", explanation="", candidate_id=f"c{i}")
            for i in range(5)
        ]
        result = filter_refined_candidates(
            candidates, eval_ctx, forest, "original", max_candidates=5
        )
        assert 5 == len(result)

    def test_over_limit_trims(self, eval_ctx):
        """
        More candidates than max_candidates trims the result.
        """
        forest = CandidateForest()
        candidates = [
            Candidate(code=f"code{i}", explanation="", candidate_id=f"c{i}")
            for i in range(10)
        ]
        result = filter_refined_candidates(
            candidates, eval_ctx, forest, "original", max_candidates=3
        )
        assert 3 == len(result)

    def test_prefers_lower_parent_runtime(self):
        """
        Candidates with faster parents rank higher.
        """
        eval_ctx = EvaluationContext()
        forest = CandidateForest()

        parent_fast = Candidate(code="pf", explanation="", candidate_id="pf")
        parent_slow = Candidate(code="ps", explanation="", candidate_id="ps")
        forest.add(parent_fast)
        forest.add(parent_slow)
        eval_ctx.record_success("pf", runtime=10.0, speedup=5.0)
        eval_ctx.record_success("ps", runtime=1000.0, speedup=0.1)

        candidates = [
            Candidate(
                code="slow_child",
                explanation="",
                candidate_id="c_slow",
                parent_id="ps",
            ),
            Candidate(
                code="fast_child",
                explanation="",
                candidate_id="c_fast",
                parent_id="pf",
            ),
        ] * 4

        for i, c in enumerate(candidates):
            candidates[i] = attrs.evolve(
                c, candidate_id=f"{c.candidate_id}_{i}"
            )

        result = filter_refined_candidates(
            candidates, eval_ctx, forest, "original", max_candidates=2
        )
        assert 2 == len(result)
        assert all("c_fast" in r.candidate_id for r in result)


class TestSelectBest:
    """Tests for select_best."""

    def test_empty_returns_none(self, eval_ctx):
        """
        No candidates returns None.
        """
        assert select_best(eval_ctx, 1000, [], []) is None

    def test_single_candidate(self, eval_ctx):
        """
        A single candidate is always selected.
        """
        eval_ctx.record_success("c1", runtime=50.0, speedup=1.0)
        assert "c1" == select_best(eval_ctx, 1000, [10], ["c1"])

    def test_picks_faster_candidate(self, eval_ctx):
        """
        The faster candidate wins when diff lengths are equal.
        """
        eval_ctx.record_success("c1", runtime=500.0, speedup=1.0)
        eval_ctx.record_success("c2", runtime=100.0, speedup=4.0)
        result = select_best(eval_ctx, 1000, [10, 10], ["c1", "c2"])
        assert "c2" == result

    def test_diff_length_contributes_to_score(self, eval_ctx):
        """
        Candidate with shorter diff and equal runtime wins when listed first.
        """
        eval_ctx.record_success("short_diff", runtime=100.0, speedup=1.0)
        eval_ctx.record_success("long_diff", runtime=100.0, speedup=1.0)
        result = select_best(
            eval_ctx, 1000, [5, 500], ["short_diff", "long_diff"]
        )
        assert "short_diff" == result

    def test_runtime_weighted_more_than_diff(self, eval_ctx):
        """
        A faster candidate wins even with a longer diff.
        """
        eval_ctx.record_success("c1", runtime=500.0, speedup=1.0)
        eval_ctx.record_success("c2", runtime=50.0, speedup=9.0)
        result = select_best(eval_ctx, 1000, [5, 100], ["c1", "c2"])
        assert "c2" == result

    def test_missing_runtime_uses_original(self, eval_ctx):
        """
        A candidate without recorded runtime uses original_runtime_ns.
        """
        eval_ctx.record_success("c1", runtime=50.0, speedup=1.0)
        result = select_best(eval_ctx, 1000, [10, 10], ["c1", "c2"])
        assert "c1" == result
