from __future__ import annotations

import ast
import logging
import re
import uuid
from pathlib import Path
from typing import TYPE_CHECKING

import libcst as cst

from codeflash_api.diff._base import DiffMethod
from codeflash_api.diff._search_replace import SearchAndReplaceDiff
from codeflash_api.diff._v4a import V4ADiff
from codeflash_api.languages.python._cst_utils import (
    find_init,
    parse_module_to_cst,
)
from codeflash_api.languages.python._markdown import (
    extract_code_block_with_context,
    group_code,
    is_multi_context,
    split_markdown_code,
    wrap_code_in_markdown,
)
from codeflash_api.languages.python._postprocess import (
    OptimizationCandidate,
    optimizations_postprocessing_pipeline,
)
from codeflash_api.optimize.schemas import OptimizeResponseItem

if TYPE_CHECKING:
    from codeflash_api.diff._base import Diff

log = logging.getLogger(__name__)

_PROMPTS_DIR = (
    Path(__file__).parent.parent / "languages" / "python" / "prompts"
)

EXPLANATION_THEN_CODE = (
    "Begin your response with a short explanation of the changes"
    " (without any title or heading like 'Explanation' or 'Changes'),"
    " then output the code."
)
FULL_CODE_PROMPT_INSTRUCTIONS = (
    "- Always provide the FULL, updated content of the artifact."
    " This means:\n"
    "    - Include ALL updated code, even if some parts are"
    " unchanged\n"
    '    - NEVER use placeholders like "// rest of the code'
    ' remains the same..." or "<- leave original code here ->"\n'
)

MARKDOWN_CONTEXT_PROMPT = (
    (_PROMPTS_DIR / "multi_file_code_format.md")
    .read_text()
    .format(EXPLANATION_THEN_CODE=EXPLANATION_THEN_CODE)
)
DEPS_CONTEXT_PROMPT = (
    _PROMPTS_DIR / "dependency_context_prompt.md"
).read_text()
INIT_OPTIMIZATION_PROMPT = (
    _PROMPTS_DIR / "init_optimization_prompt.md"
).read_text()
LINE_PROF_CONTEXT_PROMPT = (
    _PROMPTS_DIR / "lineprof_context_prompt.md"
).read_text()
RUNTIME_CONTEXT_PROMPT = (
    _PROMPTS_DIR / "runtime_context_prompt.md"
).read_text()
TEST_EXAMPLES_PROMPT = (_PROMPTS_DIR / "test_examples_prompt.md").read_text()

V4A_DIFF_FORMAT_PROMPT = """Describe the changes using the V4A diff format, enclosed within `*** Begin Patch` and `*** End Patch` markers.
# V4A Diff Format Rules:

Your entire response containing the patch MUST start with `*** Begin Patch` on a line by itself.
Your entire response containing the patch MUST end with `*** End Patch` on a line by itself.
Use the *FULL* file path, as shown to you by the user.

For each file you need to modify, start with a marker line:
    *** Update File: [path/to/file]

You ONLY update existing files, do not add or remove files.

Each file MUST appear only once in the patch.
Consolidate all changes for that file into the same block.

For `Update` actions, describe each snippet of code that needs to be changed using the following format:
1. Context lines: Include 3 lines of context *before* the change. These lines MUST start with a single space ` `.
2. Lines to remove: Precede each line to be removed with a minus sign `-`.
3. Lines to add: Precede each line to be added with a plus sign `+`.
4. Context lines: Include 3 lines of context *after* the change. These lines MUST start with a single space ` `.

Context lines MUST exactly match the existing file content, character for character, including indentation.
If a change is near the beginning or end of the file, include fewer than 3 context lines as appropriate.
If 3 lines of context is insufficient to uniquely identify the snippet, use `@@ [CLASS_OR_FUNCTION_NAME]` markers on their own lines *before* the context lines to specify the scope.
Do not include line numbers.

Only create patches for files that the user has provided.

ONLY EVER RETURN CODE IN THE SPECIFIED V4A DIFF FORMAT, followed by a short explanation of the changes.
"""

SEARCH_AND_REPLACE_FORMAT_PROMPT = """Describe the changes using the search/replace diff format with xml-style tags, like:

<replace_in_file>
<path>src/main.py</path>
<diff>
<<<<<<< SEARCH
a = 2
=======
a = 3
>>>>>>> REPLACE
</diff>
</replace_in_file>

Description: Request to replace sections of content in an existing file using SEARCH/REPLACE blocks that define exact changes to specific parts of the file.
Parameters:
- path: (required) The path of the file to modify
- diff: (required) One or more SEARCH/REPLACE blocks following this exact format:
  ```
  <<<<<<< SEARCH
  [exact content to find]
  =======
  [new content to replace with]
  >>>>>>> REPLACE
  ```
  Critical rules:
  1. SEARCH content must match the associated file section to find EXACTLY.
  2. SEARCH/REPLACE blocks will ONLY replace the first match occurrence.
  3. Keep SEARCH/REPLACE blocks concise.
  4. Special operations:
     * To move code: Use two SEARCH/REPLACE blocks
     * To delete code: Use empty REPLACE section
"""


def _humanize_ns(ns: int) -> str:
    """
    Convert nanoseconds to a human-readable string.
    """
    if ns < 1_000:
        return f"{ns}ns"
    if ns < 1_000_000:
        return f"{ns / 1_000:.1f}us"
    if ns < 1_000_000_000:
        return f"{ns / 1_000_000:.1f}ms"
    return f"{ns / 1_000_000_000:.2f}s"


def parse_python_version(version: str | None) -> tuple[int, int, int]:
    """
    Parse a version string like '3.12.9' into a tuple.
    """
    if not version:
        msg = "Python version is required"
        raise ValueError(msg)
    if len(version) > 30:
        msg = "Invalid version format"
        raise ValueError(msg)
    parts = version.split(".")
    if len(parts) != 3:
        msg = "Invalid version format"
        raise ValueError(msg)
    patch_match = re.match(r"\d+", parts[2])
    if not patch_match:
        msg = "Invalid patch version"
        raise ValueError(msg)
    major = int(parts[0])
    minor = int(parts[1])
    patch = int(patch_match.group())
    if major != 3:
        msg = "Only Python 3 is supported"
        raise ValueError(msg)
    if minor < 9 or minor > 15:
        msg = "Only Python 3.9 and above is supported"
        raise ValueError(msg)
    if patch < 0 or patch >= 100:
        msg = "Invalid version format"
        raise ValueError(msg)
    return (major, minor, patch)


def validate_trace_id(trace_id: str) -> bool:
    """
    Check that *trace_id* is a valid UUIDv4.
    """
    normalized = trace_id
    if trace_id[-4:] in ("EXP0", "EXP1"):
        normalized = trace_id[:-4] + "0000"
    try:
        uuid_obj = uuid.UUID(normalized, version=4)
        return str(uuid_obj) == normalized
    except (ValueError, AttributeError):
        return False


class CodeStrAndExplanation:
    """
    Extracted code and explanation from LLM output.
    """

    __slots__ = ("code", "explanation")

    def __init__(self, code: str, explanation: str) -> None:
        self.code = code
        self.explanation = explanation


class BaseOptimizerContext:
    """
    Base context for Python optimization.
    """

    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:
        """
        Factory: choose Single or Multi based on source format.
        """
        if is_multi_context(source_code):
            file_to_code = split_markdown_code(source_code)
            files = list(file_to_code.keys())
            if len(files) == 1 and diff_method == DiffMethod.NO_DIFF:
                file_name = files[0]
                return SingleOptimizerContext(
                    system_prompt,
                    user_prompt,
                    source_code=file_to_code[file_name],
                    file_name=file_name,
                )
            return MultiOptimizerContext(
                system_prompt,
                user_prompt,
                source_code,
                diff_method=diff_method,
            )
        return SingleOptimizerContext(system_prompt, user_prompt, source_code)

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

    def get_user_prompt(
        self,
        dependency_code: str,
        line_profiler_results: str | None,
        *,
        baseline_runtime_ns: int | None = None,
        loop_count: int | None = None,
        test_input_examples: str | None = None,
    ) -> str:
        """
        Return the formatted user prompt.
        """
        return self.base_user_prompt

    def extract_code_and_explanation_from_llm_res(
        self, content: str
    ) -> CodeStrAndExplanation:
        """
        Parse LLM output into code and explanation.
        """
        raise NotImplementedError

    def parse_and_generate_candidate_schema(
        self,
    ) -> OptimizeResponseItem | None:
        """
        Post-process extracted code into a response item.
        """
        raise NotImplementedError

    def is_valid_code(self) -> bool:
        """
        Check if extracted code is non-empty.
        """
        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[int, ...],
    ) -> None:
        """
        Validate that *code* is syntactically valid Python.
        """
        final_code = code or self.source_code
        ver = (feature_version[0], feature_version[1])
        parsed = ast.parse(final_code, feature_version=ver)
        if not parsed.body:
            msg = "Source code is empty or invalid"
            raise SyntaxError(msg)


class SingleOptimizerContext(BaseOptimizerContext):
    """
    Context for single-file optimizations.
    """

    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:
        """
        Format system prompt with full-code instructions.
        """
        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,
        *,
        baseline_runtime_ns: int | None = None,
        loop_count: int | None = None,
        test_input_examples: str | None = None,
    ) -> str:
        """
        Build user prompt with context sections.
        """
        markdown_source_code = wrap_code_in_markdown(self.source_code)
        runtime_part = ""
        if baseline_runtime_ns is not None and loop_count is not None:
            runtime_part = RUNTIME_CONTEXT_PROMPT.format(
                baseline_runtime_ns=baseline_runtime_ns,
                baseline_runtime_human=_humanize_ns(baseline_runtime_ns),
                loop_count=loop_count,
            )
        test_examples_part = ""
        if test_input_examples:
            test_examples_part = TEST_EXAMPLES_PROMPT.format(
                test_input_examples=test_input_examples
            )
        has_init = find_init(ast.parse(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 has_init else '')}"
            f"{runtime_part}"
            f"{LINE_PROF_CONTEXT_PROMPT.format(line_profiler_results=line_profiler_results) if line_profiler_results else ''}"
            f"{test_examples_part}"
        )

    def extract_code_and_explanation_from_llm_res(
        self, content: str
    ) -> CodeStrAndExplanation:
        """
        Extract code block and explanation from LLM response.
        """
        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,
    ) -> OptimizeResponseItem | None:
        """
        Post-process extracted code into a response item.
        """
        if self.extracted_code_and_expl is None:
            msg = "Call extract_code_and_explanation_from_llm_res first"
            raise AttributeError(msg)
        extracted = self.extracted_code_and_expl
        if extracted.code == "":
            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 = optimizations_postprocessing_pipeline(
                original_cst_module,
                [
                    OptimizationCandidate(
                        cst_module=cst_module,
                        id=op_id,
                        explanation=extracted.explanation,
                    )
                ],
            )

            if len(postprocessed_list) == 0:
                return None

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

        except (
            ValueError,
            cst.ParserSyntaxError,
        ):
            log.warning("Error parsing optimization result", exc_info=True)
            return None

    def is_valid_code(self) -> bool:
        """
        Check if extracted code is valid.
        """
        return super().is_valid_code()

    def validate_and_parse_source_code(
        self,
        code: str,
        feature_version: tuple[int, ...],
    ) -> None:
        """
        Validate single-file source code.
        """
        compile(self.source_code, "source_code", "exec")
        super().validate_and_parse_source_code(
            self.source_code, feature_version
        )


class MultiOptimizerContext(BaseOptimizerContext):
    """
    Context for multi-file optimizations.
    """

    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:
        """
        Format system prompt based on diff method.
        """
        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:
            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,
        *,
        baseline_runtime_ns: int | None = None,
        loop_count: int | None = None,
        test_input_examples: str | None = None,
    ) -> str:
        """
        Build user prompt for multi-file context.
        """
        has_init = any(
            find_init(ast.parse(code))
            for code in self.original_file_to_code.values()
        )
        runtime_part = ""
        if baseline_runtime_ns is not None and loop_count is not None:
            runtime_part = RUNTIME_CONTEXT_PROMPT.format(
                baseline_runtime_ns=baseline_runtime_ns,
                baseline_runtime_human=_humanize_ns(baseline_runtime_ns),
                loop_count=loop_count,
            )
        test_examples_part = ""
        if test_input_examples:
            test_examples_part = TEST_EXAMPLES_PROMPT.format(
                test_input_examples=test_input_examples
            )
        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"{runtime_part}"
            f"{LINE_PROF_CONTEXT_PROMPT.format(line_profiler_results=line_profiler_results) if line_profiler_results else ''}"
            f"{test_examples_part}"
        )

    def extract_code_and_explanation_from_llm_res(
        self, content: str
    ) -> CodeStrAndExplanation:
        """
        Extract code from LLM response based on diff method.
        """
        markdown_code = ""
        explanation = ""
        if self.diff_method == DiffMethod.NO_DIFF:
            explanation = content.split("```python:", maxsplit=1)[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)
            patch_end_marker = "*** End Patch"
            index_of_end = content.rfind(patch_end_marker)
            explanation = content[
                index_of_end + len(patch_end_marker) :
            ].strip()
        elif self.diff_method == DiffMethod.SEARCH_AND_REPLACE:
            diff = SearchAndReplaceDiff(
                content=content,
                source_code=self.original_file_to_code,
            )
            new_file_to_code = diff.run()
            end_tag = "</replace_in_file>"
            index_of_end = content.rfind(end_tag)
            explanation = content[index_of_end + len(end_tag) :].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,
    ) -> OptimizeResponseItem | None:
        """
        Post-process each extracted file individually.
        """
        if self.extracted_code_and_expl is None:
            msg = "Call extract_code_and_explanation_from_llm_res first"
            raise AttributeError(msg)

        extracted = self.extracted_code_and_expl
        if extracted.code == "":
            return None

        original_code_files = dict(self.original_file_to_code.items())
        extracted_code_files = dict(
            split_markdown_code(extracted.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
            )
        )

        valid_extracted_files = {
            f: c
            for f, c in extracted_code_files.items()
            if f in original_code_files
        }
        if not valid_extracted_files:
            log.warning(
                "No matching files. Original: %s, Extracted: %s",
                list(original_code_files.keys()),
                list(extracted_code_files.keys()),
            )
            return None

        optimization_ignored = True

        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 = OptimizationCandidate(
                    cst_module=new_cst_module,
                    id=f"{original_file}:post-processing",
                    explanation=extracted.explanation,
                )
                postprocessed_list = optimizations_postprocessing_pipeline(
                    original_cst_module,
                    [code_and_explanation],
                )
                if len(postprocessed_list) == 0:
                    new_post_processed_code[original_file] = new_code
                    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,
                cst.ParserSyntaxError,
            ):
                log.warning(
                    "Error processing file %s",
                    original_file,
                    exc_info=True,
                )
                return None

        if optimization_ignored:
            return None
        return OptimizeResponseItem(
            optimization_id=op_id,
            explanation=new_explanation,
            source_code=group_code(new_post_processed_code),
        )

    def is_valid_code(self) -> bool:
        """
        Validate extracted files are subset of originals.
        """
        if self.extracted_code_and_expl is None:
            return False
        extracted_files = set(
            split_markdown_code(self.extracted_code_and_expl.code).keys()
        )

        invalid_files = extracted_files - self._original_files_set
        if invalid_files:
            log.warning(
                "LLM returned files not in original: %s",
                invalid_files,
            )
            return False

        return super().is_valid_code()

    def validate_and_parse_source_code(
        self,
        code: str,
        feature_version: tuple[int, ...],
    ) -> None:
        """
        Validate each file in multi-file source code.
        """
        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)
