from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import libcst as cst
from pydantic import ValidationError

from core.languages.python.cst_utils import parse_module_to_cst
from aiservice.common.markdown_utils import split_markdown_code, wrap_code_in_markdown
from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
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


@dataclass()
class RefinementContextData:
    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"  # python, javascript, typescript
    language_version: str | None = None


##########################################################################################
#                               BaseRefinerContext                                      #
##########################################################################################
class BaseRefinerContext:
    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 | None, base_system_prompt: str, base_user_prompt: str
    ) -> BaseRefinerContext:
        if ctx_data is None:
            msg = "ctx_data cannot be None"
            raise ValueError(msg)
        if is_multi_context(ctx_data.optimized_source_code):
            return MultiRefinerContext(ctx_data, base_system_prompt, base_user_prompt)
        # older cli versions
        return SingleRefinerContext(ctx_data, base_system_prompt, base_user_prompt)

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

    def get_user_prompt(self) -> str:
        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:
        raise NotImplementedError

    def is_valid_refinement(self, new_refined_code: str) -> bool:
        stripped_code = new_refined_code.strip()
        if not stripped_code:
            return False

        if stripped_code == self.data.optimized_source_code.strip():
            return False

        if self.data.language in ("javascript", "typescript"):
            if self.data.language == "typescript":
                valid, _ = validate_typescript_syntax(stripped_code)
            else:
                valid, _ = validate_javascript_syntax(stripped_code)
            return bool(valid)

        try:
            parse_module_to_cst(new_refined_code)
            return True
        except cst.ParserSyntaxError:
            return False

    def validate_code_syntax(self, code: str) -> None:
        """Validate code syntax based on language. Raises exception if invalid."""
        raise NotImplementedError

    # Keep for backward compatibility
    def validate_python_module(self, code: str) -> None:
        return self.validate_code_syntax(code)


##########################################################################################
#                               SingleRefinerContext                                     #
##########################################################################################
class SingleRefinerContext(BaseRefinerContext):
    def __init__(self, ctx_data: RefinementContextData, base_system_prompt: str, base_user_prompt: str) -> None:
        super().__init__(ctx_data, base_system_prompt, base_user_prompt)

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

    def get_user_prompt(self) -> str:
        # Use appropriate language for code fence
        lang = self.data.language if self.data.language else "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:
        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 is_valid_refinement(self, new_refined_code: str) -> bool:
        return super().is_valid_refinement(new_refined_code)

    def validate_code_syntax(self, code: str) -> None:
        """Validate code syntax based on language."""
        if self.data.language in ("javascript", "typescript"):
            if self.data.language == "typescript":
                valid, _ = validate_typescript_syntax(code)
                lang_name = "TypeScript"
            else:
                valid, _ = validate_javascript_syntax(code)
                lang_name = "JavaScript"
            if not valid:
                raise ValueError(f"Invalid {lang_name} syntax")
            return

        # Python validation using libcst
        try:
            cst_module = parse_module_to_cst(code)
            CodeAndExplanation(cst_module, "")
        except (ValueError, ValidationError, cst.ParserSyntaxError):
            raise


##########################################################################################
#                               MultiRefinerContext                                      #
##########################################################################################
class MultiRefinerContext(BaseRefinerContext):
    def __init__(self, ctx_data: RefinementContextData, base_system_prompt: str, base_user_prompt: str) -> None:
        super().__init__(ctx_data, base_system_prompt, base_user_prompt)

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

    def apply_patches_to_optimized_code(self, llm_res: str) -> str:
        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:
        if is_markdown_structure_changed(new_refined_code, self.data.optimized_source_code):
            return False
        valid = True
        for code in split_markdown_code(new_refined_code).values():
            if not super().is_valid_refinement(code):
                valid = False
                break
        return valid

    def validate_code_syntax(self, code: str) -> None:
        """Validate code syntax based on language."""
        # For JavaScript/TypeScript, use the appropriate validator
        if self.data.language in ("javascript", "typescript"):
            if self.data.language == "typescript":
                valid, _ = validate_typescript_syntax(code)
                lang_name = "TypeScript"
            else:
                valid, _ = validate_javascript_syntax(code)
                lang_name = "JavaScript"
            if not valid:
                raise ValueError(f"Invalid {lang_name} syntax")
            return

        # Python validation using libcst
        for code in split_markdown_code(self.data.optimized_source_code).values():
            try:
                cst_module = parse_module_to_cst(code)
                CodeAndExplanation(cst_module, "")
            except (ValueError, ValidationError, cst.ParserSyntaxError):
                raise
