"""Tests for early candidate deduplication via dedup_candidates."""

from __future__ import annotations

from codeflash_core import Candidate, EvaluationContext, dedup_candidates
from codeflash_python.analysis._normalizer import normalize_python_code


def make_candidate(
    code: str,
    cid: str | None = None,
    source: str = "optimize",
) -> Candidate:
    """Build a Candidate with sensible defaults."""
    return Candidate(
        code=code,
        explanation="test",
        candidate_id=cid or f"opt-{id(code)}",
        source=source,
    )


def normalize_fn(source: str) -> str:
    """Normalize, keeping the candidate on failure."""
    try:
        return normalize_python_code(source, remove_docstrings=True)
    except Exception:
        return source


ORIGINAL_CODE = "def foo(x):\n    return x + 1\n"
ORIGINAL_FLAT = f"# file: test.py\n{ORIGINAL_CODE}"

# Normalizes identically to ORIGINAL_CODE (docstring and comment stripped)
IDENTICAL_TO_ORIGINAL = (
    'def foo(x):\n    """Docstring."""\n    # comment\n    return x + 1\n'
)

# Different from original
CANDIDATE_A = "def foo(x):\n    return x + 2\n"
CANDIDATE_B = "def foo(x):\n    return x * 2\n"
CANDIDATE_C = "def foo(x):\n    return x << 1\n"

# Normalizes identically to CANDIDATE_A (added comment stripped by normalizer)
CANDIDATE_A_DUP = "def foo(x):\n    # optimized\n    return x + 2\n"

NORMALIZED_ORIGINAL = normalize_fn(ORIGINAL_CODE.strip())


def run_dedup(
    candidates: list[Candidate],
    *,
    seen: set[str] | None = None,
    cross_batch: dict[str, dict[str, object]] | None = None,
) -> list[Candidate]:
    """Convenience wrapper around dedup_candidates."""
    return dedup_candidates(
        candidates,
        normalize_fn=normalize_fn,
        original_normalized=NORMALIZED_ORIGINAL,
        seen=seen,
        cross_batch=cross_batch,
    )


class TestDedup:
    """Tests for the dedup_candidates pipeline function."""

    def test_unique_candidates_pass_through(self) -> None:
        """All unique candidates survive dedup."""
        candidates = [
            make_candidate(CANDIDATE_A, "opt-a"),
            make_candidate(CANDIDATE_B, "opt-b"),
            make_candidate(CANDIDATE_C, "opt-c"),
        ]
        unique = run_dedup(candidates)
        assert 3 == len(unique)

    def test_identical_to_original_removed(self) -> None:
        """Candidates that normalize to the original are removed."""
        candidates = [
            make_candidate(IDENTICAL_TO_ORIGINAL, "opt-dup-orig"),
            make_candidate(CANDIDATE_A, "opt-a"),
        ]
        unique = run_dedup(candidates)
        assert 1 == len(unique)

    def test_intra_batch_duplicates_removed(self) -> None:
        """Duplicates within a single batch are removed."""
        candidates = [
            make_candidate(CANDIDATE_A, "opt-a1"),
            make_candidate(CANDIDATE_A_DUP, "opt-a2"),
            make_candidate(CANDIDATE_B, "opt-b"),
        ]
        unique = run_dedup(candidates)
        assert 2 == len(unique)

    def test_cross_batch_duplicates_removed(self) -> None:
        """Candidates matching a prior batch (via cross_batch) are removed."""
        eval_ctx = EvaluationContext()
        normalized_a = normalize_fn(CANDIDATE_A.strip())
        eval_ctx.register_new(
            normalized_a,
            "opt-prev",
            CANDIDATE_A,
            ORIGINAL_FLAT,
        )
        eval_ctx.record_success("opt-prev", runtime=1000.0, speedup=2.0)

        new_candidates = [
            make_candidate(CANDIDATE_A_DUP, "opt-new-dup"),
            make_candidate(CANDIDATE_B, "opt-b"),
        ]
        unique = run_dedup(
            new_candidates,
            cross_batch=eval_ctx.code_to_id,
        )
        # Only CANDIDATE_B survives (A_DUP is a cross-batch dup)
        assert 1 == len(unique)

    def test_empty_list(self) -> None:
        """Empty input returns empty output."""
        unique = run_dedup([])
        assert 0 == len(unique)

    def test_all_duplicates_of_original(self) -> None:
        """All candidates identical to the original are removed."""
        candidates = [
            make_candidate(IDENTICAL_TO_ORIGINAL, "opt-1"),
            make_candidate(ORIGINAL_CODE, "opt-2"),
        ]
        unique = run_dedup(candidates)
        assert 0 == len(unique)

    def test_mixed_removal_types(self) -> None:
        """Intra-batch dups, original dups, and cross-batch dups are all removed."""
        eval_ctx = EvaluationContext()
        normalized_c = normalize_fn(CANDIDATE_C.strip())
        eval_ctx.register_new(
            normalized_c,
            "opt-prev-c",
            CANDIDATE_C,
            ORIGINAL_FLAT,
        )
        eval_ctx.record_success("opt-prev-c", runtime=500.0, speedup=3.0)

        candidates = [
            make_candidate(
                IDENTICAL_TO_ORIGINAL, "opt-orig"
            ),  # identical to original
            make_candidate(CANDIDATE_A, "opt-a1"),  # unique
            make_candidate(
                CANDIDATE_A_DUP, "opt-a2"
            ),  # intra-batch dup of opt-a1
            make_candidate(CANDIDATE_C, "opt-c-dup"),  # cross-batch dup
            make_candidate(CANDIDATE_B, "opt-b"),  # unique
        ]
        unique = run_dedup(
            candidates,
            cross_batch=eval_ctx.code_to_id,
        )
        # Only CANDIDATE_A and CANDIDATE_B should survive
        assert 2 == len(unique)

    def test_dedup_with_shared_seen_set(self) -> None:
        """A shared seen set prevents duplicates across sequential batches."""
        seen: set[str] = set()

        first_batch = [make_candidate(CANDIDATE_A, "opt-a")]
        unique_first = run_dedup(first_batch, seen=seen)
        assert 1 == len(unique_first)

        # Second batch reuses the same seen set
        second_batch = [
            make_candidate(CANDIDATE_B, "opt-b"),
            make_candidate(CANDIDATE_A_DUP, "opt-a-lp"),  # dup of first batch
        ]
        unique_second = run_dedup(second_batch, seen=seen)
        assert 1 == len(unique_second)
        assert "opt-b" == unique_second[0].candidate_id


class TestEvaluationContext:
    """Direct tests for EvaluationContext register/handle_duplicate with original_flat_code."""

    def test_register_new_stores_diff_len(self) -> None:
        """register_new records candidate_id, shorter_code, and diff_len."""
        eval_ctx = EvaluationContext()
        normalized = normalize_fn(CANDIDATE_A.strip())

        eval_ctx.register_new(
            normalized,
            "opt-a",
            CANDIDATE_A,
            ORIGINAL_FLAT,
        )

        entry = eval_ctx.code_to_id[normalized]
        assert "opt-a" == entry["candidate_id"]
        assert CANDIDATE_A == entry["shorter_code"]
        assert isinstance(entry["diff_len"], int)
        assert entry["diff_len"] > 0

    def test_handle_duplicate_copies_all_results(self) -> None:
        """handle_duplicate copies speedup, runtime, correctness, and line profiler."""
        eval_ctx = EvaluationContext()
        normalized = normalize_fn(CANDIDATE_A.strip())

        eval_ctx.register_new(
            normalized,
            "opt-first",
            CANDIDATE_A,
            ORIGINAL_FLAT,
        )
        eval_ctx.record_success("opt-first", runtime=1234.0, speedup=2.5)
        eval_ctx.record_line_profile("opt-first", "line profiler output")

        eval_ctx.handle_duplicate(
            "opt-dup",
            normalized,
            ORIGINAL_FLAT,
            CANDIDATE_A_DUP,
        )

        assert 2.5 == eval_ctx.speedup_ratios["opt-dup"]
        assert 1234.0 == eval_ctx.optimized_runtimes["opt-dup"]
        assert eval_ctx.is_correct["opt-dup"] is True
        assert (
            "line profiler output" == eval_ctx.line_profiler_results["opt-dup"]
        )

    def test_handle_duplicate_copies_failed_results(self) -> None:
        """handle_duplicate copies failure state for failed candidates."""
        eval_ctx = EvaluationContext()
        normalized = normalize_fn(CANDIDATE_A.strip())

        eval_ctx.register_new(
            normalized,
            "opt-first",
            CANDIDATE_A,
            ORIGINAL_FLAT,
        )
        eval_ctx.record_failed("opt-first")

        eval_ctx.handle_duplicate(
            "opt-dup",
            normalized,
            ORIGINAL_FLAT,
            CANDIDATE_A_DUP,
        )

        assert eval_ctx.speedup_ratios["opt-dup"] is None
        assert eval_ctx.optimized_runtimes["opt-dup"] is None
        assert eval_ctx.is_correct["opt-dup"] is False

    def test_handle_duplicate_tracks_shorter_source(self) -> None:
        """When a duplicate has a shorter diff, it replaces shorter_code."""
        eval_ctx = EvaluationContext()
        longer_code = "def foo(x):\n    # this comment makes it longer\n    # and this one too\n    return x + 2\n"
        normalized = normalize_fn(longer_code.strip())

        eval_ctx.register_new(
            normalized,
            "opt-long",
            longer_code,
            ORIGINAL_FLAT,
        )
        eval_ctx.record_success("opt-long", runtime=500.0, speedup=3.0)
        original_diff_len = eval_ctx.code_to_id[normalized]["diff_len"]

        # Duplicate with shorter code (same normalized form)
        eval_ctx.handle_duplicate(
            "opt-short",
            normalized,
            ORIGINAL_FLAT,
            CANDIDATE_A,
        )
        new_diff_len = eval_ctx.code_to_id[normalized]["diff_len"]

        # Shorter code should have replaced the longer one
        assert new_diff_len <= original_diff_len
        assert CANDIDATE_A == eval_ctx.code_to_id[normalized]["shorter_code"]

    def test_handle_duplicate_keeps_shorter_when_new_is_longer(
        self,
    ) -> None:
        """When a duplicate has a longer diff, the original shorter_code is kept."""
        eval_ctx = EvaluationContext()
        normalized = normalize_fn(CANDIDATE_A.strip())

        eval_ctx.register_new(
            normalized,
            "opt-short",
            CANDIDATE_A,
            ORIGINAL_FLAT,
        )
        eval_ctx.record_success("opt-short", runtime=500.0, speedup=3.0)

        longer_code = "def foo(x):\n    # this comment makes it longer\n    # and this one too\n    return x + 2\n"
        eval_ctx.handle_duplicate(
            "opt-long",
            normalized,
            ORIGINAL_FLAT,
            longer_code,
        )

        assert CANDIDATE_A == eval_ctx.code_to_id[normalized]["shorter_code"]
