from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import attrs
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,
    is_multi_context,
    split_markdown_code,
    wrap_code_in_markdown,
)

if TYPE_CHECKING:
    from codeflash_api.diff._base import Diff

log = logging.getLogger(__name__)


@attrs.frozen
class RefinementContextData:
    """
    Immutable data bag for a single refinement item.
    """

    original_source_code: str
    original_line_profiler_results: str | None
    original_code_runtime: str
    read_only_dependency_code: str | None
    optimized_source_code: str
    optimized_line_profiler_results: str | None
    optimized_code_runtime: str
    speedup: str
    optimized_explanation: str = ""
    python_version: str | None = None
    function_references: str | None = None
    language: str = "python"
    language_version: str | None = None


class BaseRefinerContext:
    """
    Base context for refinement prompt assembly and patch application.
    """

    def __init__(
        self,
        ctx_data: RefinementContextData,
        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

    @staticmethod
    def get_dynamic_context(
        ctx_data: RefinementContextData,
        base_system_prompt: str,
        base_user_prompt: str,
    ) -> BaseRefinerContext:
        """
        Factory: choose Single or Multi based on source format.
        """
        if is_multi_context(ctx_data.optimized_source_code):
            return MultiRefinerContext(
                ctx_data, base_system_prompt, base_user_prompt
            )
        return SingleRefinerContext(
            ctx_data, base_system_prompt, base_user_prompt
        )

    def get_system_prompt(self) -> str:
        """
        Return the system prompt.
        """
        return self.base_system_prompt

    def get_user_prompt(self) -> str:
        """
        Return the formatted user prompt with all placeholders filled.
        """
        return self.base_user_prompt.format(
            original_source_code=self.data.original_source_code,
            original_line_profiler_results=(
                self.data.original_line_profiler_results
                or "[No profiler results available]"
            ),
            optimized_source_code=self.data.optimized_source_code,
            optimized_line_profiler_results=(
                self.data.optimized_line_profiler_results
                or "[No profiler results available]"
            ),
            optimized_explanation=self.data.optimized_explanation,
            original_code_runtime=self.data.original_code_runtime,
            optimized_code_runtime=self.data.optimized_code_runtime,
            speedup=self.data.speedup,
            read_only_dependency_code=(
                self.data.read_only_dependency_code
                or "[No read only code present]"
            ),
            python_version=(self.data.python_version or "Not Available"),
            function_references=(
                self.data.function_references or "Not Available"
            ),
        )

    def apply_patches_to_optimized_code(self, llm_res: str) -> str:
        """
        Apply SEARCH/REPLACE patches from *llm_res* to the optimized code.
        """
        raise NotImplementedError

    def is_valid_refinement(self, new_refined_code: str) -> bool:
        """
        Check that *new_refined_code* is non-empty, different from the
        optimized source, and syntactically valid.
        """
        stripped = new_refined_code.strip()
        if not stripped:
            return False
        if stripped == self.data.optimized_source_code.strip():
            return False
        try:
            parse_module_to_cst(new_refined_code)
        except cst.ParserSyntaxError:
            return False
        return True

    def validate_code_syntax(self, code: str) -> None:
        """
        Validate *code* parses as a Python module via libcst.
        """
        raise NotImplementedError


class SingleRefinerContext(BaseRefinerContext):
    """
    Context for single-file refinements.
    """

    def get_user_prompt(self) -> str:
        """
        Format user prompt with markdown-wrapped code blocks.
        """
        lang = self.data.language or "python"
        return self.base_user_prompt.format(
            original_source_code=wrap_code_in_markdown(
                self.data.original_source_code, lang
            ),
            original_line_profiler_results=(
                self.data.original_line_profiler_results
                or "[No profiler results available]"
            ),
            optimized_source_code=wrap_code_in_markdown(
                self.data.optimized_source_code, lang
            ),
            optimized_line_profiler_results=(
                self.data.optimized_line_profiler_results
                or "[No profiler results available]"
            ),
            optimized_explanation=self.data.optimized_explanation,
            original_code_runtime=self.data.original_code_runtime,
            optimized_code_runtime=self.data.optimized_code_runtime,
            speedup=self.data.speedup,
            read_only_dependency_code=(
                self.data.read_only_dependency_code
                or "[No read only code present]"
            ),
            python_version=(self.data.python_version or "Not Available"),
            function_references=(
                self.data.function_references or "Not Available"
            ),
        )

    def apply_patches_to_optimized_code(self, llm_res: str) -> str:
        """
        Apply patches to single-file optimized code.
        """
        diff: Diff = SearchAndReplaceDiff(
            content=llm_res,
            source_code={"file": self.data.optimized_source_code},
            match_files_when_having_single_patch=False,
        )
        file_to_code = diff.run()
        return file_to_code["file"]

    def validate_code_syntax(self, code: str) -> None:
        """
        Validate single-file Python syntax via libcst.
        """
        parse_module_to_cst(code)


class MultiRefinerContext(BaseRefinerContext):
    """
    Context for multi-file refinements.
    """

    def apply_patches_to_optimized_code(self, llm_res: str) -> str:
        """
        Apply patches to multi-file optimized code.
        """
        file_to_code = split_markdown_code(self.data.optimized_source_code)
        diff: Diff = SearchAndReplaceDiff(
            content=llm_res, source_code=file_to_code
        )
        file_to_code = diff.run()
        return group_code(file_to_code)

    def is_valid_refinement(self, new_refined_code: str) -> bool:
        """
        Validate that file structure is preserved and each file parses.
        """
        old_files = set(
            split_markdown_code(self.data.optimized_source_code).keys()
        )
        new_files = set(split_markdown_code(new_refined_code).keys())
        if old_files != new_files:
            return False
        for code in split_markdown_code(new_refined_code).values():
            if not super().is_valid_refinement(code):
                return False
        return True

    def validate_code_syntax(self, code: str) -> None:
        """
        Validate each file in a multi-file source via libcst.
        """
        for file_code in split_markdown_code(
            self.data.optimized_source_code
        ).values():
            parse_module_to_cst(file_code)
