from __future__ import annotations

import ast
import textwrap

import pytest

from codeflash_python.verification._ranking import (
    CandidateEvaluationContext,
    create_rank_dictionary_compact,
    diff_length,
    normalize_code,
    normalize_node,
    select_best_candidate,
)


class TestNormalizeNode:
    """normalize_node AST docstring and import stripping."""

    def test_strips_module_docstring(self) -> None:
        """Module-level docstring is removed from body."""
        src = textwrap.dedent("""\
            \"\"\"Module docstring.\"\"\"
            x = 1
        """)
        tree = ast.parse(src)
        result = normalize_node(tree)

        unparsed = ast.unparse(result)
        assert "Module docstring" not in unparsed
        assert "x = 1" in unparsed

    def test_strips_function_docstring(self) -> None:
        """Function-level docstring is removed."""
        src = textwrap.dedent("""\
            def foo():
                \"\"\"Foo docstring.\"\"\"
                return 1
        """)
        tree = ast.parse(src)
        result = normalize_node(tree)

        unparsed = ast.unparse(result)
        assert "Foo docstring" not in unparsed
        assert "return 1" in unparsed

    def test_strips_class_docstring(self) -> None:
        """Class-level docstring is removed."""
        src = textwrap.dedent("""\
            class Bar:
                \"\"\"Bar docstring.\"\"\"
                x = 1
        """)
        tree = ast.parse(src)
        result = normalize_node(tree)

        unparsed = ast.unparse(result)
        assert "Bar docstring" not in unparsed
        assert "x = 1" in unparsed

    def test_strips_import_statements(self) -> None:
        """Import and import-from statements are removed."""
        src = textwrap.dedent("""\
            import os
            from sys import argv
            x = 1
        """)
        tree = ast.parse(src)
        result = normalize_node(tree)

        unparsed = ast.unparse(result)
        assert "import os" not in unparsed
        assert "from sys" not in unparsed
        assert "x = 1" in unparsed

    def test_leaves_other_statements_intact(self) -> None:
        """Assignments and expressions survive normalization."""
        src = textwrap.dedent("""\
            x = 1
            y = x + 2
        """)
        tree = ast.parse(src)
        result = normalize_node(tree)

        unparsed = ast.unparse(result)
        assert "x = 1" in unparsed
        assert "y = x + 2" in unparsed

    def test_nested_function_inside_class(self) -> None:
        """Docstrings are stripped from nested structures."""
        src = textwrap.dedent("""\
            class Outer:
                \"\"\"Outer doc.\"\"\"
                def method(self):
                    \"\"\"Method doc.\"\"\"
                    return 42
        """)
        tree = ast.parse(src)
        result = normalize_node(tree)

        unparsed = ast.unparse(result)
        assert "Outer doc" not in unparsed
        assert "Method doc" not in unparsed
        assert "return 42" in unparsed


class TestNormalizeCode:
    """normalize_code string-level AST normalization."""

    def test_strips_docstrings_and_imports(self) -> None:
        """Code differing only in docstrings and imports normalizes equal."""
        code_a = textwrap.dedent("""\
            \"\"\"Module A.\"\"\"
            import os
            def foo():
                \"\"\"Foo A.\"\"\"
                return 1
        """)
        code_b = textwrap.dedent("""\
            \"\"\"Module B.\"\"\"
            import sys
            def foo():
                \"\"\"Foo B.\"\"\"
                return 1
        """)

        assert normalize_code(code_a) == normalize_code(code_b)

    def test_different_variable_names_not_equal(self) -> None:
        """Different variable names produce different output."""
        code_a = "x = 1"
        code_b = "y = 1"

        assert normalize_code(code_a) != normalize_code(code_b)

    def test_invalid_code_returns_raw_string(self) -> None:
        """Invalid Python code falls back to the raw string."""
        normalize_code.cache_clear()
        result = normalize_code("def (broken:")

        assert "def (broken:" == result


class TestDiffLength:
    """diff_length unified diff character count."""

    def test_identical_strings_return_zero(self) -> None:
        """Identical strings have zero diff length."""
        assert 0 == diff_length("hello\n", "hello\n")

    def test_different_strings_return_positive(self) -> None:
        """Different strings produce a positive diff length."""
        result = diff_length("hello\n", "world\n")

        assert result > 0

    def test_empty_vs_nonempty_returns_positive(self) -> None:
        """Empty string versus non-empty produces a positive diff."""
        result = diff_length("", "something\n")

        assert result > 0


class TestCreateRankDictionaryCompact:
    """create_rank_dictionary_compact index-to-rank mapping."""

    def test_ascending_list(self) -> None:
        """Already sorted list: rank equals index."""
        result = create_rank_dictionary_compact([10, 20, 30])

        assert {0: 0, 1: 1, 2: 2} == result

    def test_descending_list(self) -> None:
        """Descending list: ranks are reversed."""
        result = create_rank_dictionary_compact([30, 20, 10])

        assert {0: 2, 1: 1, 2: 0} == result

    def test_single_element(self) -> None:
        """Single element always gets rank 0."""
        result = create_rank_dictionary_compact([42])

        assert {0: 0} == result

    def test_ties_all_indices_present(self) -> None:
        """Tied values still produce a valid mapping with all indices."""
        result = create_rank_dictionary_compact([5, 5, 5])

        assert {0, 1, 2} == set(result.keys())
        assert {0, 1, 2} == set(result.values())


class TestCandidateEvaluationContext:
    """CandidateEvaluationContext mutable tracking state."""

    def test_record_failed_candidate(self) -> None:
        """Failed candidate sets runtime=None, correct=False, speedup=None."""
        ctx = CandidateEvaluationContext()
        ctx.record_failed_candidate("opt-1")

        assert ctx.optimized_runtimes["opt-1"] is None
        assert ctx.is_correct["opt-1"] is False
        assert ctx.speedup_ratios["opt-1"] is None

    def test_record_successful_candidate(self) -> None:
        """Successful candidate stores runtime, correct=True, and speedup."""
        ctx = CandidateEvaluationContext()
        ctx.record_successful_candidate("opt-2", runtime=500.0, speedup=2.0)

        assert 500.0 == ctx.optimized_runtimes["opt-2"]
        assert ctx.is_correct["opt-2"] is True
        assert 2.0 == ctx.speedup_ratios["opt-2"]

    def test_record_line_profiler_result(self) -> None:
        """Line profiler result is stored under the optimization id."""
        ctx = CandidateEvaluationContext()
        ctx.record_line_profiler_result("opt-3", "profile data")

        assert "profile data" == ctx.optimized_line_profiler_results["opt-3"]

    def test_register_new_candidate(self) -> None:
        """New candidate creates an entry in ast_code_to_id."""
        ctx = CandidateEvaluationContext()
        ctx.register_new_candidate(
            normalized_code="x = 1",
            optimization_id="opt-4",
            source_code_flat="x = 1\n",
            original_flat_code="y = 2\n",
        )

        assert "x = 1" in ctx.ast_code_to_id
        assert "opt-4" == ctx.ast_code_to_id["x = 1"]["optimization_id"]

    def test_handle_duplicate_candidate(self) -> None:
        """Duplicate candidate copies results from prior evaluation."""
        ctx = CandidateEvaluationContext()
        ctx.register_new_candidate(
            normalized_code="x = 1",
            optimization_id="opt-first",
            source_code_flat="x = 1\n",
            original_flat_code="y = 2\n",
        )
        ctx.record_successful_candidate(
            "opt-first", runtime=100.0, speedup=3.0
        )

        ctx.handle_duplicate_candidate(
            optimization_id="opt-dup",
            normalized_code="x = 1",
            candidate_source_code_flat="x = 1\n",
            original_flat_code="y = 2\n",
        )

        assert 3.0 == ctx.speedup_ratios["opt-dup"]
        assert ctx.is_correct["opt-dup"] is True
        assert 100.0 == ctx.optimized_runtimes["opt-dup"]

    def test_handle_duplicate_updates_shorter_source(self) -> None:
        """Duplicate with shorter diff updates the stored source code."""
        ctx = CandidateEvaluationContext()
        original = "y = 2\nz = 3\n"
        ctx.register_new_candidate(
            normalized_code="x = 1",
            optimization_id="opt-long",
            source_code_flat="x = 1\nextra_long_line = True\n",
            original_flat_code=original,
        )
        ctx.record_successful_candidate("opt-long", runtime=100.0, speedup=2.0)

        shorter_source = "x = 1\n"
        ctx.handle_duplicate_candidate(
            optimization_id="opt-short",
            normalized_code="x = 1",
            candidate_source_code_flat=shorter_source,
            original_flat_code=original,
        )

        assert (
            shorter_source
            == ctx.ast_code_to_id["x = 1"]["shorter_source_code"]
        )

    def test_get_speedup_ratio_unknown_id(self) -> None:
        """Unknown optimization id returns None for speedup ratio."""
        ctx = CandidateEvaluationContext()

        assert ctx.get_speedup_ratio("nonexistent") is None

    def test_get_optimized_runtime_unknown_id(self) -> None:
        """Unknown optimization id returns None for runtime."""
        ctx = CandidateEvaluationContext()

        assert ctx.get_optimized_runtime("nonexistent") is None


class TestSelectBestCandidate:
    """select_best_candidate combined ranking selection."""

    def test_empty_optimization_ids_returns_none(self) -> None:
        """No candidates returns None."""
        ctx = CandidateEvaluationContext()

        result = select_best_candidate(
            eval_ctx=ctx,
            original_runtime_ns=1000,
            diff_lengths=[],
            optimization_ids=[],
        )

        assert result is None

    def test_single_candidate_returns_its_id(self) -> None:
        """Single candidate is always selected."""
        ctx = CandidateEvaluationContext()
        ctx.record_successful_candidate("opt-a", runtime=500.0, speedup=2.0)

        result = select_best_candidate(
            eval_ctx=ctx,
            original_runtime_ns=1000,
            diff_lengths=[10],
            optimization_ids=["opt-a"],
        )

        assert "opt-a" == result

    def test_selects_best_combined_ranking(self) -> None:
        """Candidate with best combined diff+runtime ranking wins."""
        ctx = CandidateEvaluationContext()
        ctx.record_successful_candidate("opt-a", runtime=500.0, speedup=2.0)
        ctx.record_successful_candidate("opt-b", runtime=100.0, speedup=10.0)

        result = select_best_candidate(
            eval_ctx=ctx,
            original_runtime_ns=1000,
            diff_lengths=[10, 5],
            optimization_ids=["opt-a", "opt-b"],
        )

        # opt-b has both shorter diff (5 < 10) and faster runtime (100 < 500)
        assert "opt-b" == result

    def test_tradeoff_diff_vs_runtime(self) -> None:
        """Short diff vs fast runtime tradeoff is resolved by ranking."""
        ctx = CandidateEvaluationContext()
        # opt-a: short diff (5), slow runtime (900)
        ctx.record_successful_candidate("opt-a", runtime=900.0, speedup=1.1)
        # opt-b: long diff (100), fast runtime (50)
        ctx.record_successful_candidate("opt-b", runtime=50.0, speedup=20.0)

        result = select_best_candidate(
            eval_ctx=ctx,
            original_runtime_ns=1000,
            diff_lengths=[5, 100],
            optimization_ids=["opt-a", "opt-b"],
        )

        # Both have combined rank of 1:
        # opt-a: diff_rank=0 + runtime_rank=1
        # opt-b: diff_rank=1 + runtime_rank=0
        # Tie-break via min(); either is acceptable.
        assert result in {"opt-a", "opt-b"}
