import ast
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING

import libcst as cst
import sentry_sdk
from pydantic import ValidationError

from core.languages.python.cst_utils import find_init, parse_module_to_cst
from aiservice.common.markdown_utils import extract_code_block_with_context, split_markdown_code, wrap_code_in_markdown
from aiservice.env_specific import debug_log_sensitive_data
from core.languages.python.optimizer.context_utils.constants import (
    DEPS_CONTEXT_PROMPT,
    EXPLANATION_THEN_CODE,
    FULL_CODE_PROMPT_INSTRUCTIONS,
    INIT_OPTIMIZATION_PROMPT,
    LINE_PROF_CONTEXT_PROMPT,
    MARKDOWN_CONTEXT_PROMPT,
)
from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
from core.languages.python.optimizer.diff_patches_utils.diff import (
    SEARCH_AND_REPLACE_FORMAT_PROMPT,
    V4A_DIFF_FORMAT_PROMPT,
    DiffMethod,
)
from core.languages.python.optimizer.diff_patches_utils.search_and_replace import SearchAndReplaceDiff
from core.languages.python.optimizer.diff_patches_utils.v4a_diff import V4ADiff
from core.languages.python.optimizer.models import CodeExplanationAndID
from core.languages.python.optimizer.postprocess import optimizations_postprocessing_pipeline
from core.shared.context_helpers import group_code
from core.shared.optimizer_schemas import OptimizeResponseItemSchema

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


@dataclass
class CodeStrAndExplanation:
    code: str
    explanation: str


##########################################################################################
#                               BaseOptimizerContext [PYTHON ONLY]                       #
##########################################################################################
class BaseOptimizerContext:
    def __init__(self, base_system_prompt: str, base_user_prompt: str, source_code: str) -> None:
        self.base_system_prompt = base_system_prompt
        self.base_user_prompt = base_user_prompt
        self.source_code = source_code
        self.extracted_code_and_expl: CodeStrAndExplanation | None = None
        self.code_and_explanation_before_post_processing: dict[str, CodeStrAndExplanation] = {}

    @staticmethod
    def get_dynamic_context(
        system_prompt: str, user_prompt: str, source_code: str, diff_method: DiffMethod = DiffMethod.NO_DIFF
    ) -> "BaseOptimizerContext":
        if is_multi_context(source_code):
            file_to_code = split_markdown_code(source_code)
            files = list(file_to_code.keys())
            # diff methods uses file names so the single context won't work for it
            if len(files) == 1 and diff_method == DiffMethod.NO_DIFF:
                # multi-file context with only a single file, we can use the old single context
                file_name = files[0]
                return SingleOptimizerContext(
                    system_prompt, user_prompt, source_code=file_to_code[file_name], file_name=file_name
                )
            # only multi optimizer context can accept a diff method
            return MultiOptimizerContext(system_prompt, user_prompt, source_code, diff_method=diff_method)
        # older cli versions
        return SingleOptimizerContext(system_prompt, user_prompt, source_code)

    def get_system_prompt(self, python_version_str: str) -> str:  # noqa: ARG002
        return self.base_system_prompt

    def get_user_prompt(self, dependency_code: str, line_profiler_results: str | None) -> str:  # noqa: ARG002
        return self.base_user_prompt

    def extract_code_and_explanation_from_llm_res(self, content: str) -> CodeStrAndExplanation:
        raise NotImplementedError

    def parse_and_generate_candidate_schema(self) -> OptimizeResponseItemSchema | None:
        raise NotImplementedError

    def is_valid_code(self) -> bool:
        if self.extracted_code_and_expl is None:
            return False
        code = self.extracted_code_and_expl.code
        return code is not None and code.strip() != ""

    def validate_and_parse_source_code(self, code: str, feature_version: tuple) -> None:
        final_code = code or self.source_code
        parsed = ast.parse(final_code, feature_version=feature_version)
        if not parsed.body:
            msg = "Source code is empty or invalid"
            raise SyntaxError(msg)


##########################################################################################
#                               SingleOptimizerContext                                   #
##########################################################################################
class SingleOptimizerContext(BaseOptimizerContext):
    def __init__(
        self, base_system_prompt: str, base_user_prompt: str, source_code: str, file_name: str | None = None
    ) -> None:
        super().__init__(base_system_prompt, base_user_prompt, source_code)
        self.file_name = file_name

    def get_system_prompt(self, python_version_str: str) -> str:
        return (
            self.base_system_prompt.format(
                python_version_str=python_version_str, critical_instructions=FULL_CODE_PROMPT_INSTRUCTIONS
            )
            + "\n"
            + EXPLANATION_THEN_CODE
        )

    def get_user_prompt(self, dependency_code: str, line_profiler_results: str | None) -> str:
        markdown_source_code = wrap_code_in_markdown(self.source_code)
        return (
            f"{DEPS_CONTEXT_PROMPT.format(dependency_code=dependency_code) if dependency_code else ''}"
            f"{self.base_user_prompt.format(source_code=markdown_source_code, init_optimization_prompt=INIT_OPTIMIZATION_PROMPT if find_init(ast.parse(self.source_code)) else '')}"
            f"{LINE_PROF_CONTEXT_PROMPT.format(line_profiler_results=line_profiler_results) if line_profiler_results else ''}"
        )

    def extract_code_and_explanation_from_llm_res(self, content: str) -> CodeStrAndExplanation:
        code = ""
        explanation = ""
        if result := extract_code_block_with_context(content, language="python"):
            explanation, code, _ = result
        extracted = CodeStrAndExplanation(code=code, explanation=explanation)
        self.extracted_code_and_expl = extracted
        return extracted

    def parse_and_generate_candidate_schema(self) -> OptimizeResponseItemSchema | None:
        if self.extracted_code_and_expl is None:
            msg = "extracted_code_and_expl cannot be None, did you forget to call extract_code_and_explanation_from_llm_res"
            raise AttributeError(msg)
        extracted = self.extracted_code_and_expl
        if extracted.code == "":
            # couldn't extract the code from llm response
            return None
        try:
            op_id = str(uuid.uuid4())
            cst_module = parse_module_to_cst(extracted.code)
            self.code_and_explanation_before_post_processing[op_id] = CodeStrAndExplanation(
                code=cst_module.code, explanation=extracted.explanation
            )
            original_cst_module = parse_module_to_cst(self.source_code)
            postprocessed_list: list[CodeExplanationAndID] = optimizations_postprocessing_pipeline(
                original_cst_module,
                [CodeExplanationAndID(cst_module=cst_module, id=op_id, explanation=extracted.explanation)],
            )

            if len(postprocessed_list) == 0:
                return None

            postprocessed = postprocessed_list[0]
            return OptimizeResponseItemSchema(
                explanation=postprocessed.explanation,
                optimization_id=postprocessed.id,
                source_code=group_code({self.file_name: postprocessed.cst_module.code})
                if self.file_name is not None
                else postprocessed.cst_module.code,
            )

        except (ValueError, ValidationError, cst.ParserSyntaxError) as e:
            # log exception with sentry
            sentry_sdk.capture_exception(e)
            debug_log_sensitive_data(f"error for source:\n{extracted.code}")
            debug_log_sensitive_data(f"Traceback: {e}")
            return None

    def is_valid_code(self) -> bool:
        return super().is_valid_code()

    def validate_and_parse_source_code(self, code: str, feature_version: tuple) -> None:
        compile(self.source_code, "source_code", "exec")
        super().validate_and_parse_source_code(self.source_code, feature_version)


##########################################################################################
#                               MultiOptimizerContext                                    #
##########################################################################################
class MultiOptimizerContext(BaseOptimizerContext):
    def __init__(
        self, base_system_prompt: str, base_user_prompt: str, source_code: str, diff_method: DiffMethod
    ) -> None:
        self.diff_method = diff_method
        super().__init__(base_system_prompt, base_user_prompt, source_code)
        self.original_file_to_code = split_markdown_code(source_code)

        self._original_files_set = set(self.original_file_to_code.keys())

    def get_system_prompt(self, python_version_str: str) -> str:
        # critical instructions should be bullet points, short and direct, further instructions should be a full paragraph which is added after the main system prompt
        critical_instructions = ""
        code_format_instructions = ""
        if self.diff_method == DiffMethod.V4A:
            code_format_instructions = V4A_DIFF_FORMAT_PROMPT
            critical_instructions = "Begin your response with the diff followed by a short explanation of the changes (without any title or heading like 'Explanation' or 'Changes')."
        elif self.diff_method == DiffMethod.SEARCH_AND_REPLACE:
            code_format_instructions = SEARCH_AND_REPLACE_FORMAT_PROMPT
            critical_instructions = "Begin your response with the diff followed by a short explanation of the changes (without any title or heading like 'Explanation' or 'Changes')."
        elif self.diff_method == DiffMethod.NO_DIFF:
            # we still use the markdown context if we won't use any diff format
            code_format_instructions = MARKDOWN_CONTEXT_PROMPT
            critical_instructions = FULL_CODE_PROMPT_INSTRUCTIONS
        return (
            self.base_system_prompt.format(
                python_version_str=python_version_str, critical_instructions=critical_instructions
            )
            + "\n"
            + code_format_instructions
        )

    def get_user_prompt(self, dependency_code: str, line_profiler_results: str | None) -> str:
        has_init = any(find_init(ast.parse(code)) for code in self.original_file_to_code.values())
        return (
            f"{DEPS_CONTEXT_PROMPT.format(dependency_code=dependency_code) if dependency_code else ''}"
            f"{self.base_user_prompt.format(source_code=self.source_code, init_optimization_prompt=INIT_OPTIMIZATION_PROMPT if has_init else '')}"
            f"{LINE_PROF_CONTEXT_PROMPT.format(line_profiler_results=line_profiler_results) if line_profiler_results else ''}"
        )

    def extract_code_and_explanation_from_llm_res(self, content: str) -> CodeStrAndExplanation:
        markdown_code = ""
        explanation = ""
        if self.diff_method == DiffMethod.NO_DIFF:
            # simple markdown context is returned
            # the explanation is everything before the first code block
            explanation = content.split("```python:")[0].strip()
            markdown_code = group_code(split_markdown_code(content))
        elif self.diff_method == DiffMethod.V4A:
            diff: Diff = V4ADiff(content=content, source_code=self.original_file_to_code)
            new_file_to_code = diff.run()
            markdown_code = group_code(new_file_to_code)
            # explanation is the content from the end of the patch until the end of the string
            patch_end_marker = "*** End Patch"
            index_of_end_of_patch = content.rfind(patch_end_marker)
            explanation = content[index_of_end_of_patch + len(patch_end_marker) :].strip()
        elif self.diff_method == DiffMethod.SEARCH_AND_REPLACE:
            diff: Diff = SearchAndReplaceDiff(content=content, source_code=self.original_file_to_code)
            new_file_to_code = diff.run()
            index_of_end_of_patch = content.rfind("</replace_in_file>")
            explanation = content[index_of_end_of_patch + len("</replace_in_file>") :].strip()
            markdown_code = group_code(new_file_to_code)

        result = CodeStrAndExplanation(code=markdown_code, explanation=explanation)
        self.extracted_code_and_expl = result
        return result

    def parse_and_generate_candidate_schema(self) -> OptimizeResponseItemSchema | None:
        if self.extracted_code_and_expl is None:
            msg = "extracted_code_and_expl cannot be None, did you forget to call extract_code_and_explanation_from_llm_res"
            raise AttributeError(msg)

        extracted = self.extracted_code_and_expl

        markdown_code = extracted.code
        if markdown_code == "":
            return None

        original_code_files = dict(self.original_file_to_code.items())
        extracted_code_files = dict(split_markdown_code(markdown_code).items())

        op_id = str(uuid.uuid4())
        new_post_processed_code: dict[str, str] = {}
        new_explanation: str = ""

        self.code_and_explanation_before_post_processing[op_id] = CodeStrAndExplanation(
            code=extracted.code, explanation=extracted.explanation
        )

        # Check if any extracted files exist in the original context
        # LLMs may only return files they actually modified
        valid_extracted_files = {f: c for f, c in extracted_code_files.items() if f in original_code_files}
        if not valid_extracted_files:
            debug_log_sensitive_data(
                f"No matching files found. Original: {list(original_code_files.keys())}, "
                f"Extracted: {list(extracted_code_files.keys())}"
            )
            return None

        # Log when LLM returns fewer files (common with large contexts)
        if len(valid_extracted_files) < len(original_code_files):
            debug_log_sensitive_data(
                f"LLM returned {len(valid_extracted_files)} of {len(original_code_files)} files. "
                f"Processing returned files: {list(valid_extracted_files.keys())}"
            )

        # optimization is ignored when **all** files in the context are rejected by postprocessing pipeline
        optimization_ignored = True

        # post process each file that was returned by the LLM
        for original_file, new_code in valid_extracted_files.items():
            original_code = original_code_files[original_file]
            try:
                new_cst_module = parse_module_to_cst(new_code)
                original_cst_module = parse_module_to_cst(original_code)

                code_and_explanation = CodeExplanationAndID(
                    cst_module=new_cst_module, id=f"{original_file}:post-processing", explanation=extracted.explanation
                )
                postprocessed_list: list[CodeExplanationAndID] = optimizations_postprocessing_pipeline(
                    original_cst_module, [code_and_explanation]
                )
                if len(postprocessed_list) == 0:
                    # rejected by postprocessing pipeline
                    debug_log_sensitive_data(f"postprocessed_list is empty for file: {original_file}")
                    debug_log_sensitive_data(f"original_code: {original_code}")
                    debug_log_sensitive_data(f"new_code: {new_code}")
                    new_post_processed_code[original_file] = (
                        new_code  # still save it, in case other files are not removed by post-processing
                    )
                    continue

                optimization_ignored = False
                postprocessed = postprocessed_list[0]

                if new_explanation == "":
                    new_explanation = postprocessed.explanation

                new_post_processed_code[original_file] = postprocessed.cst_module.code

            except (ValueError, ValidationError, cst.ParserSyntaxError) as e:
                # log exception with sentry
                sentry_sdk.capture_exception(e)
                debug_log_sensitive_data(f"error for source:\n{new_code}")
                debug_log_sensitive_data(f"Traceback: {e}")
                return None
        return (
            OptimizeResponseItemSchema(
                optimization_id=op_id, explanation=new_explanation, source_code=group_code(new_post_processed_code)
            )
            if not optimization_ignored
            else None
        )

    def is_valid_code(self) -> bool:
        if self.extracted_code_and_expl is None:
            return False
        extracted_files = set(split_markdown_code(self.extracted_code_and_expl.code).keys())

        # All extracted files must be in the original set (no new files allowed)
        invalid_files = extracted_files - self._original_files_set
        if invalid_files:
            debug_log_sensitive_data(
                f"LLM returned files not in original context: {invalid_files}. "
                f"Original: {self._original_files_set}, Extracted: {extracted_files}"
            )
            return False

        return super().is_valid_code()

    def validate_and_parse_source_code(self, code: str, feature_version: tuple) -> None:
        final_code = code or self.source_code
        code_blocks = split_markdown_code(final_code).values()
        for code_block in code_blocks:
            compile(code_block, "source_code", "exec")
            super().validate_and_parse_source_code(code_block, feature_version)
