import logging
from collections import defaultdict
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, Any

import libcst as cst
import sentry_sdk
from ninja import Field, Schema
from pydantic import ValidationError

from aiservice.common.markdown_utils import split_markdown_code
from core.languages.python.cst_utils import parse_module_to_cst
from core.languages.python.optimizer.diff_patches_utils.search_and_replace import SearchAndReplaceDiff
from core.languages.python.optimizer.models import CodeAndExplanation
from core.shared.context_helpers import group_code, is_markdown_structure_changed

if TYPE_CHECKING:
    from core.languages.python.optimizer.diff_patches_utils.diff import Diff


class TestDiffScope(str, Enum):
    RETURN_VALUE = "return_value"
    STDOUT = "stdout"
    DID_PASS = "did_pass"  # noqa: S105
    TIMED_OUT = "timed_out"


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


class TestDiff(Schema):
    scope: TestDiffScope
    original_value: bool | str | int | float | dict[str, Any] | list[Any] | None = None
    candidate_value: bool | str | int | float | dict[str, Any] | list[Any] | None = None
    original_pass: bool
    candidate_pass: bool
    test_src_code: str | None = None
    candidate_pytest_error: str | None = None
    original_pytest_error: str | None = None


class CodeRepairRequestSchema(Schema):
    trace_id: str
    optimization_id: str
    original_source_code: str
    modified_source_code: str
    test_diffs: list[TestDiff] = Field(..., alias="test_diffs")
    language: str = "python"  # python, javascript, typescript
    rerun_trace_id: str | None = None


@dataclass()
class CodeRepairContextData:
    original_source_code: str
    modified_source_code: str
    test_diffs: list[TestDiff]
    language: str = "python"


class CodeRepairContext:
    def __init__(self, ctx_data: CodeRepairContextData, base_system_prompt: str, base_user_prompt: str) -> None:
        self.data = ctx_data
        self.base_system_prompt = base_system_prompt
        self.base_user_prompt = base_user_prompt

    def get_system_prompt(self) -> str:
        return self.base_system_prompt

    def build_test_details(self, test_diffs: list[TestDiff]) -> str:
        sections: defaultdict[str, list[str]] = defaultdict(list)
        language = self.data.language
        test_error_label = "Pytest error" if language == "python" else "Test error"
        seen_headers: set[str] = set()

        for diff in test_diffs:
            try:
                key = diff.test_src_code or ""
                if key not in seen_headers:
                    seen_headers.add(key)
                    header_lines = [
                        f"**Test source:**",
                        f"```{language}",
                        f"{diff.test_src_code or 'Not available'}",
                        f"```",
                    ]
                    if diff.original_pytest_error:
                        header_lines.append(f"{test_error_label} (original): {diff.original_pytest_error}")
                    if diff.candidate_pytest_error:
                        header_lines.append(f"{test_error_label} (optimized): {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 != TestDiffScope.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 as e:
                logging.exception("Some issue in parsing test diffs")
                sentry_sdk.capture_exception(e)

        parts = []
        for key, lines in sections.items():
            parts.append("\n".join(lines))
        return "\n\n---\n\n".join(parts)

    def get_user_prompt(self) -> str:
        return self.base_user_prompt.format(
            original_source_code=self.data.original_source_code,
            modified_source_code=self.data.modified_source_code,
            test_details=self.build_test_details(self.data.test_diffs),
        )

    def apply_patches_to_optimized_code(self, llm_res: str) -> str:
        file_to_code = split_markdown_code(self.data.modified_source_code, self.data.language)
        raw_input = not file_to_code
        if raw_input:
            file_to_code = {"file.py": self.data.modified_source_code}
        diff: Diff = SearchAndReplaceDiff(content=llm_res, 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, self.data.language)

    def is_valid(self, new_refined_code: str) -> bool:
        code_blocks = split_markdown_code(new_refined_code, self.data.language)
        if code_blocks:
            if is_markdown_structure_changed(new_refined_code, self.data.modified_source_code, self.data.language):
                return False
            code_items = code_blocks.values()
        else:
            code_items = [new_refined_code]
        for code in code_items:
            stripped_code = code.strip()
            if not stripped_code:
                return False
            if self.data.language == "python":
                try:
                    parse_module_to_cst(code)
                except cst.ParserSyntaxError:
                    return False
        return True

    def validate_module(self) -> None:
        """Validate the module syntax based on language."""
        # Skip validation for non-Python languages for now
        # TODO: have some way to validate the syntax of the code for other languages like js & ts
        if self.data.language != "python":
            return
        for _code in split_markdown_code(self.data.modified_source_code).values():
            try:
                cst_module = parse_module_to_cst(_code)
                CodeAndExplanation(cst_module, "")
            except (ValueError, ValidationError, cst.ParserSyntaxError):  # noqa: TRY203
                raise
