from __future__ import annotations

import logging
from collections import defaultdict

import libcst as cst

from codeflash_api.diff._search_replace import (
    SearchAndReplaceDiff,
)
from codeflash_api.languages.python._cst_utils import (
    parse_module_to_cst,
)
from codeflash_api.languages.python._markdown import (
    group_code,
    split_markdown_code,
)
from codeflash_api.repair.schemas import (
    BehaviorDiff,
    BehaviorDiffScope,
)

log = logging.getLogger(__name__)

SCOPE_DESCRIPTIONS: dict[BehaviorDiffScope, str] = {
    BehaviorDiffScope.RETURN_VALUE: (
        "The function returned a different value in the"
        " optimized code compared to the original."
    ),
    BehaviorDiffScope.STDOUT: (
        "The output printed to stdout is different in the"
        " optimized code compared to the original."
    ),
    BehaviorDiffScope.DID_PASS: (
        "The test passed in one version but failed in the"
        " other (a change in pass/fail behavior)."
    ),
}


def build_test_details(
    test_diffs: list[BehaviorDiff],
    language: str = "python",
) -> str:
    """
    Format *test_diffs* into a markdown section for the user prompt.
    """
    sections: defaultdict[str, list[str]] = defaultdict(list)
    test_error_label = "Pytest error" if language == "python" else "Test error"
    seen_headers: set[str] = set()

    for diff in test_diffs:
        _format_single_diff(
            diff, sections, seen_headers, test_error_label, language
        )

    parts = ["\n".join(lines) for lines in sections.values()]
    return "\n\n---\n\n".join(parts)


def _format_single_diff(
    diff: BehaviorDiff,
    sections: defaultdict[str, list[str]],
    seen_headers: set[str],
    test_error_label: str,
    language: str,
) -> None:
    """
    Format one *diff* and append it to *sections*.
    """
    try:
        key = diff.test_src_code or ""
        if key not in seen_headers:
            seen_headers.add(key)
            header_lines = [
                "**Test source:**",
                f"```{language}",
                f"{diff.test_src_code or 'Not available'}",
                "```",
            ]
            if diff.original_pytest_error:
                header_lines.append(
                    f"{test_error_label} (original):"
                    f" {diff.original_pytest_error}"
                )
            if diff.candidate_pytest_error:
                header_lines.append(
                    f"{test_error_label} (optimized):"
                    f" {diff.candidate_pytest_error}"
                )
            sections[key].append("\n".join(header_lines))

        scope_desc = SCOPE_DESCRIPTIONS.get(diff.scope, diff.scope.value)
        detail_lines = [f"- {scope_desc}"]
        if diff.scope != BehaviorDiffScope.DID_PASS:
            detail_lines.append(f"  Expected: {diff.original_value!r}")
            detail_lines.append(f"  Got:      {diff.candidate_value!r}")
        else:
            orig_status = "Passed" if diff.original_pass else "Failed"
            opt_status = "Passed" if diff.candidate_pass else "Failed"
            detail_lines.append(
                f"  Original: {orig_status} → Optimized: {opt_status}"
            )
        sections[key].append("\n".join(detail_lines))
    except Exception:
        log.exception("Error parsing test diff")


def build_user_prompt(
    template: str,
    original_source_code: str,
    modified_source_code: str,
    test_diffs: list[BehaviorDiff],
    language: str = "python",
) -> str:
    """
    Format the user prompt template with repair context.
    """
    test_details = build_test_details(test_diffs, language)
    return template.format(
        original_source_code=original_source_code,
        modified_source_code=modified_source_code,
        test_details=test_details,
    )


def apply_patches_to_optimized_code(
    llm_response: str,
    modified_source_code: str,
    language: str = "python",
) -> str:
    """
    Apply SEARCH/REPLACE patches from *llm_response* to *modified_source_code*.
    """
    file_to_code = split_markdown_code(modified_source_code, language)
    raw_input = not file_to_code
    if raw_input:
        file_to_code = {"file.py": modified_source_code}

    diff = SearchAndReplaceDiff(content=llm_response, source_code=file_to_code)
    file_to_code = diff.run()

    if raw_input:
        return file_to_code.get("file.py", "")
    return group_code(file_to_code, language)


def _is_markdown_structure_changed(
    old: str, new: str, language: str = "python"
) -> bool:
    """
    Return True if the set of file paths changed between markdown blocks.
    """
    old_files = split_markdown_code(old, language).keys()
    new_files = split_markdown_code(new, language).keys()
    return old_files != new_files


def is_valid_repair(
    repaired_code: str,
    modified_source_code: str,
    language: str = "python",
) -> bool:
    """
    Validate that *repaired_code* is syntactically valid.
    """
    code_blocks = split_markdown_code(repaired_code, language)
    if code_blocks:
        if _is_markdown_structure_changed(
            repaired_code, modified_source_code, language
        ):
            return False
        code_items = list(code_blocks.values())
    else:
        code_items = [repaired_code]

    for code in code_items:
        stripped = code.strip()
        if not stripped:
            return False
        if language == "python":
            try:
                parse_module_to_cst(code)
            except cst.ParserSyntaxError:
                return False
    return True
