from __future__ import annotations

import ast
import io
import logging
import re
import tokenize
from difflib import SequenceMatcher
from typing import TYPE_CHECKING, cast

import libcst as cst
import sentry_sdk
from libcst import BaseStatement, CSTTransformer, CSTVisitor, Expr, IndentedBlock, SimpleStatementLine, SimpleString

from aiservice.common_utils import safe_isort
from core.languages.python.cst_utils import compare_unparsed_ast_to_source, parse_module_to_cst, unparse_parse_source
from core.languages.python.optimizer.models import CodeExplanationAndID
from core.languages.python.testgen.postprocessing.add_missing_imports import add_future_annotations_import

if TYPE_CHECKING:
    from libcst import FunctionDef


def deduplicate_optimizations(
    _original_module: cst.Module, optimized_code_and_explanations: list[CodeExplanationAndID]
) -> list[CodeExplanationAndID]:
    """Remove optimizations that have equivalent code.

    This function assumes that the 'cst_module' attribute of CodeExplanationAndID objects is a
    libcst.Module instance of the optimized code, and that two CodeExplanationAndID objects are
    considered duplicates if their 'cst_module' attributes produce the same code when converted
    to string.


    Args:
    ----
        _original_module: The original source code CST module (unused).
        optimized_code_and_explanations (List[CodeExplanationAndID]): A list of CodeExplanationAndID
            objects representing the optimized code and their explanations.

    Returns:
    -------
        List[CodeExplanationAndID]: A list of CodeExplanationAndID objects with duplicates removed.

    """
    seen_asts: set[str] = set()
    unique_optimizations = []
    for optimization in optimized_code_and_explanations:
        normalized = ast.unparse(ast.parse(optimization.cst_module.code))
        if normalized not in seen_asts:
            seen_asts.add(normalized)
            unique_optimizations.append(optimization)
    return unique_optimizations


def equality_check(
    original_module: cst.Module,
    optimized_code_and_explanations: list[CodeExplanationAndID],
    *,
    original_code: str | None = None,
) -> list[CodeExplanationAndID]:
    original_source_code = original_code if original_code is not None else original_module.code
    try:
        original_source_ast = unparse_parse_source(original_source_code)
    except Exception:
        return [ce for ce in optimized_code_and_explanations if ce.cst_module.code != original_source_code]
    filtered_optimizations = []
    for ce in optimized_code_and_explanations:
        try:
            if not compare_unparsed_ast_to_source(original_source_ast, ce.cst_module.code):
                filtered_optimizations.append(ce)
        except Exception:
            if ce.cst_module.code != original_source_code:
                filtered_optimizations.append(ce)
    return filtered_optimizations


explanation_sub_patterns = [
    (
        re.compile(
            r"\sHere (is|are) (the )?((code|optimization)|\S+\scode|(optimized|improved) versions? of (the code|these functions))(:|.)",
            re.IGNORECASE,
        ),
        "",
    ),
    (re.compile(r"^```(.*?)```", re.MULTILINE | re.DOTALL), ""),
    (re.compile(r", as follows:"), "."),
    (re.compile(r":\n"), ".\n"),
]


def cleanup_explanations(
    _original_module: cst.Module, optimized_code_and_explanations: list[CodeExplanationAndID]
) -> list[CodeExplanationAndID]:
    new_optimized_code_and_explanations = []

    for ce in optimized_code_and_explanations:
        cleaned_up_explanation = ce.explanation
        for pattern, repl in explanation_sub_patterns:
            cleaned_up_explanation = pattern.sub(repl, cleaned_up_explanation)

        new_optimized_code_and_explanations.append(
            CodeExplanationAndID(cst_module=ce.cst_module, explanation=cleaned_up_explanation, id=ce.id)
        )

    return new_optimized_code_and_explanations


class DocstringVisitor(CSTVisitor):
    def __init__(self) -> None:
        self.original_docstrings: dict[str, str] = {}
        self.class_name: str | None = None

    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
        self.class_name = node.name.value
        docstring = node.get_docstring(clean=False)
        if docstring:
            self.original_docstrings[self.class_name] = docstring
        return True

    def leave_ClassDef(self, original_node: cst.ClassDef) -> None:  # noqa: ARG002
        self.class_name = None

    def visit_FunctionDef(self, node: FunctionDef) -> bool:
        function_name = node.name.value
        qualified_name = f"{self.class_name}.{function_name}" if self.class_name else function_name
        docstring = node.get_docstring(clean=False)
        if docstring:
            self.original_docstrings[qualified_name] = docstring
        return True


class DocstringTransformer(CSTTransformer):
    def __init__(self, original_docstrings: dict[str, str]) -> None:
        self.original_docstrings = original_docstrings
        self.class_name: str | None = None

    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
        self.class_name = node.name.value
        return True

    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:  # noqa: ARG002
        original_docstring = self.original_docstrings.get(self.class_name) if self.class_name else None
        if original_docstring:
            if not updated_node.get_docstring(clean=False):
                new_body: list[BaseStatement] = [
                    SimpleStatementLine(body=[Expr(value=SimpleString(f'"""{original_docstring}"""'))]),
                    *cast("list[BaseStatement]", list(updated_node.body.body)),
                ]
                updated_node = updated_node.with_changes(body=IndentedBlock(body=new_body))
            else:
                new_body = [
                    SimpleStatementLine(body=[Expr(value=SimpleString(f'"""{original_docstring}"""'))]),
                    *cast("list[BaseStatement]", list(updated_node.body.body[1:])),
                ]
                updated_node = updated_node.with_changes(body=IndentedBlock(body=new_body))
        self.class_name = None
        return updated_node

    def leave_FunctionDef(self, original_node: FunctionDef, updated_node: FunctionDef) -> FunctionDef:  # noqa: ARG002
        function_name = updated_node.name.value
        qualified_name = f"{self.class_name}.{function_name}" if self.class_name else function_name
        original_docstring = self.original_docstrings.get(qualified_name)
        if original_docstring:
            if not updated_node.get_docstring(clean=False):
                new_body: list[BaseStatement] = [
                    SimpleStatementLine(body=[Expr(value=SimpleString(f'"""{original_docstring}"""'))]),
                    *cast("list[BaseStatement]", list(updated_node.body.body)),
                ]
                updated_node = updated_node.with_changes(body=IndentedBlock(body=new_body))
            else:
                new_body = [
                    SimpleStatementLine(body=[Expr(value=SimpleString(f'"""{original_docstring}"""'))]),
                    *cast("list[BaseStatement]", list(updated_node.body.body[1:])),
                ]
                updated_node = updated_node.with_changes(body=IndentedBlock(body=new_body))
        return updated_node


def fix_missing_docstring(
    original_module: cst.Module, optimized_code_and_explanations: list[CodeExplanationAndID]
) -> list[CodeExplanationAndID]:
    visitor = DocstringVisitor()
    try:
        original_module.visit(visitor)
    except Exception:
        return optimized_code_and_explanations
    original_docstrings = visitor.original_docstrings
    transformer = DocstringTransformer(original_docstrings)
    returns = []
    for ce in optimized_code_and_explanations:
        try:
            returns.append(
                CodeExplanationAndID(cst_module=ce.cst_module.visit(transformer), explanation=ce.explanation, id=ce.id)
            )
        except Exception as e:
            sentry_sdk.capture_exception(e)
            returns.append(ce)
    return returns


def dedup_and_sort_imports(
    _original_module: cst.Module, optimized_code_and_explanations: list[CodeExplanationAndID]
) -> list[CodeExplanationAndID]:
    new_optimized_code_and_explanations = []
    for ce in optimized_code_and_explanations:
        try:
            original_code = ce.cst_module.code
            sorted_code = safe_isort(original_code, disregard_skip=True)
        except Exception:
            new_optimized_code_and_explanations.append(ce)
            continue
        # Skip re-parse if isort didn't change anything
        if sorted_code == original_code:
            new_optimized_code_and_explanations.append(ce)
        else:
            new_optimized_code_and_explanations.append(
                CodeExplanationAndID(cst_module=parse_module_to_cst(sorted_code), explanation=ce.explanation, id=ce.id)
            )

    return new_optimized_code_and_explanations


class EllipsisContainingCodeVisitor(CSTVisitor):
    def __init__(self) -> None:
        self.ellipsis_containing_code = False

    def visit_Ellipsis(self, node: cst.Ellipsis) -> bool:  # noqa: ARG002
        self.ellipsis_containing_code = True
        return False


def filter_ellipsis_containing_code(
    original_module: cst.Module, optimized_code_and_explanations: list[CodeExplanationAndID]
) -> list[CodeExplanationAndID]:
    new_optimized_code_and_explanations = []
    original_visitor = EllipsisContainingCodeVisitor()
    original_module.visit(original_visitor)
    if original_visitor.ellipsis_containing_code:
        # don't check for ellipsis containing optimized code if the original code contains ellipsis
        return optimized_code_and_explanations
    for ce in optimized_code_and_explanations:
        visitor = EllipsisContainingCodeVisitor()
        ce.cst_module.visit(visitor)
        if not visitor.ellipsis_containing_code:
            new_optimized_code_and_explanations.append(ce)
    return new_optimized_code_and_explanations


def _strip_comments_from_code(code: str) -> str:
    """Remove all comments from Python code while preserving strings and their content.

    Uses tokenize to properly distinguish between actual comments and # symbols
    inside strings (including multi-line strings).

    Args:
    ----
        code: Python source code as a string

    Returns:
    -------
        The same code with all comments removed, preserving string content

    """
    # add comment here
    try:
        lines = code.splitlines(keepends=True)
        tokens = tokenize.generate_tokens(io.StringIO(code).readline)

        # Build a per-line map of comment ranges to remove
        comments_by_line: list[list[tuple[int, int]]] = [[] for _ in range(len(lines))]
        for token in tokens:
            if token.type == tokenize.COMMENT:
                line_idx = token.start[0] - 1
                if 0 <= line_idx < len(lines):
                    comments_by_line[line_idx].append((token.start[1], token.end[1]))

        # Remove comments from the code
        result_lines = []
        for line_num, line in enumerate(lines, start=1):
            # Find all comments on this line
            line_comments = comments_by_line[line_num - 1]

            if line_comments:
                # Remove the comments from this line
                # Sort by start position (rightmost first) to remove from right to left
                line_comments.sort(reverse=True)
                for start_col, end_col in line_comments:
                    # Remove the comment and any trailing whitespace before it
                    line = line[:start_col].rstrip() + line[end_col:]

                # Add back the line ending if it was removed
                if not line.endswith("\n") and lines[line_num - 1].endswith("\n"):
                    line += "\n"

            result_lines.append(line)

        return "".join(result_lines)
    except (tokenize.TokenError, IndentationError):
        # If tokenization fails, return the code as-is
        # This can happen with incomplete or malformed code
        return code


def clean_extraneous_comments(
    original_module: cst.Module,
    optimized_module: cst.Module,
    *,
    orig_code: str | None = None,
    orig_code_stripped: str | None = None,
) -> cst.Module:
    """Clean extraneous comments from optimized code using difflib.

    Uses diff-based approach on code (without comments) to identify which lines
    actually changed, then removes comments from unchanged lines.

    Args:
    ----
        original_module: The original CST module.
        optimized_module: The optimized CST module with potential extra comments.
        orig_code: Pre-computed original module code (avoids redundant codegen).
        orig_code_stripped: Pre-computed comment-stripped original code.

    Returns:
    -------
        A CST module with extraneous comments removed.

    """
    try:
        # Get line-by-line representation, reusing pre-computed strings when available
        if orig_code is None:
            orig_code = original_module.code
        opt_code_full = optimized_module.code

        orig_lines = orig_code.splitlines(keepends=True)
        opt_lines = opt_code_full.splitlines(keepends=True)

        # Strip comments from entire code to identify code changes
        # This properly handles # symbols inside strings
        if orig_code_stripped is None:
            orig_code_stripped = _strip_comments_from_code(orig_code)
        opt_code_stripped = _strip_comments_from_code(opt_code_full)

        # Split stripped versions into lines
        orig_code_only = orig_code_stripped.splitlines(keepends=True)
        opt_code_only = opt_code_stripped.splitlines(keepends=True)

        # Filter out comment-only/blank lines for comparison
        # Keep track of indices of actual code lines
        orig_code_line_indices = []
        orig_code_lines_filtered = []
        for i, line in enumerate(orig_code_only):
            if line.strip():  # Has actual code
                orig_code_line_indices.append(i)
                orig_code_lines_filtered.append(line)

        opt_code_line_indices = []
        opt_code_lines_filtered = []
        for i, line in enumerate(opt_code_only):
            if line.strip():  # Has actual code
                opt_code_line_indices.append(i)
                opt_code_lines_filtered.append(line)

        # Find which lines in optimized code have actual code changes
        # Compare only the actual code lines (not comment-only lines)
        code_matcher = SequenceMatcher(None, orig_code_lines_filtered, opt_code_lines_filtered)
        code_changed_line_indices = set()  # Indices in the filtered lists

        for tag, _i1, _i2, j1, j2 in code_matcher.get_opcodes():
            if tag != "equal":
                # These are changed/inserted/deleted lines
                for j in range(j1, j2):
                    code_changed_line_indices.add(j)

        # Map back to actual line numbers in opt_lines
        code_changed_lines = set()
        for filtered_idx in code_changed_line_indices:
            if filtered_idx < len(opt_code_line_indices):
                code_changed_lines.add(opt_code_line_indices[filtered_idx])

        # Build a mapping between original and optimized lines based on code (without comments)
        # This helps us understand which original lines correspond to which optimized lines
        orig_to_opt_mapping = {}
        opt_to_orig_mapping = {}

        for tag, i1, i2, j1, j2 in code_matcher.get_opcodes():
            if tag == "equal":
                # Lines are the same (without comments)
                for offset in range(min(i2 - i1, j2 - j1)):
                    orig_idx = orig_code_line_indices[i1 + offset]
                    opt_idx = opt_code_line_indices[j1 + offset]
                    orig_to_opt_mapping[orig_idx] = opt_idx
                    opt_to_orig_mapping[opt_idx] = orig_idx

        # Now build the result: for each line in optimized code,
        # keep it if the code changed, or if it's a comment before changed code
        result_lines = []
        orig_idx = 0
        restored_orig_indices = set()  # Track which original lines have been restored

        for opt_idx, opt_line in enumerate(opt_lines):
            if opt_idx in code_changed_lines:
                # Code changed on this line - keep the optimized version with comments
                result_lines.append(opt_line)
            else:
                # Code didn't change on this line
                # Get the corresponding stripped line
                opt_code = opt_code_only[opt_idx] if opt_idx < len(opt_code_only) else ""

                # Check if this is a comment-only line that comes before OR after a code change
                is_comment_only = not opt_code.strip()
                is_near_change = False
                if is_comment_only:
                    # Check if any of the following lines have code changes
                    for check_idx in range(opt_idx + 1, len(opt_lines)):
                        if check_idx in code_changed_lines:
                            is_near_change = True
                            break
                        # Stop checking if we hit another line with actual code
                        check_code = opt_code_only[check_idx] if check_idx < len(opt_code_only) else ""
                        if check_code.strip():
                            break

                    # Also check if any of the preceding lines (looking back) have code changes
                    if not is_near_change:
                        for check_idx in range(opt_idx - 1, -1, -1):
                            if check_idx in code_changed_lines:
                                is_near_change = True
                                break
                            # Stop checking if we hit another line with actual code
                            check_code = opt_code_only[check_idx] if check_idx < len(opt_code_only) else ""
                            if check_code.strip():
                                break

                if is_comment_only and is_near_change:
                    # Comment line near a code change - keep it as-is
                    result_lines.append(opt_line)
                elif is_comment_only and not is_near_change:
                    # Comment line not near a code change - skip it
                    pass
                # Code line - check if it changed or not
                # If changed, use optimized version. If unchanged, use original version.
                elif opt_idx in code_changed_lines:
                    # Code changed - just use the optimized line as-is, no restoration
                    result_lines.append(opt_line)
                else:
                    # Code didn't change - find and use original (including any preceding comments)
                    found_orig = None
                    orig_line_idx = None
                    for orig_idx_search in range(orig_idx, len(orig_lines)):
                        orig_code = orig_code_only[orig_idx_search] if orig_idx_search < len(orig_code_only) else ""
                        if orig_code == opt_code:
                            found_orig = orig_lines[orig_idx_search]
                            orig_line_idx = orig_idx_search
                            orig_idx = orig_idx_search + 1
                            break

                    if found_orig:
                        # Check if there are comment-only/blank lines in the original that come before this line
                        # BUT ONLY restore them if this code line is UNCHANGED (not in code_changed_lines)
                        # Don't restore original comments before CHANGED code lines
                        if orig_line_idx is not None and orig_line_idx > 0 and opt_idx not in code_changed_lines:
                            # Look backwards for ALL consecutive comment-only or blank lines
                            # Collect them all, then decide which ones to restore
                            preceding_lines: list[str] = []
                            check_idx = orig_line_idx - 1

                            while check_idx >= 0:
                                check_code = orig_code_only[check_idx] if check_idx < len(orig_code_only) else ""
                                if not check_code.strip():
                                    # This is a comment-only or blank line in the original
                                    # Add it if not already restored and if it was removed from optimized
                                    if check_idx not in orig_to_opt_mapping and check_idx not in restored_orig_indices:
                                        preceding_lines.insert(0, orig_lines[check_idx])
                                        restored_orig_indices.add(check_idx)
                                    check_idx -= 1
                                else:
                                    # Hit a line with actual code - stop looking backwards
                                    break

                            # Add the restored comments/blank lines before the actual line
                            result_lines.extend(preceding_lines)

                        # Use the original line (preserves original comments or lack thereof)
                        result_lines.append(found_orig)
                        if orig_line_idx is not None:
                            restored_orig_indices.add(orig_line_idx)

                        # Also check for trailing blank/comment lines after this line that were removed
                        # BUT: only restore them if this code line is UNCHANGED and they don't come before a changed line
                        # (in that case, the new comment from optimized should be kept)
                        if (
                            orig_line_idx is not None
                            and orig_line_idx < len(orig_lines) - 1
                            and opt_idx not in code_changed_lines
                        ):
                            trailing_lines = []
                            check_idx = orig_line_idx + 1
                            found_changed_line = False

                            while check_idx < len(orig_lines):
                                check_code = orig_code_only[check_idx] if check_idx < len(orig_code_only) else ""
                                if not check_code.strip():
                                    # This is a comment-only or blank line in the original
                                    # Check if it was removed (not in the optimized version)
                                    if check_idx not in orig_to_opt_mapping and check_idx not in restored_orig_indices:
                                        trailing_lines.append(orig_lines[check_idx])
                                        restored_orig_indices.add(check_idx)
                                    check_idx += 1
                                else:
                                    # Hit a line with actual code
                                    # Check if this code line was changed in the optimized version
                                    if check_idx in orig_to_opt_mapping:
                                        next_opt_idx = orig_to_opt_mapping[check_idx]
                                        if next_opt_idx in code_changed_lines:
                                            found_changed_line = True
                                    else:
                                        # This original line is not in the mapping, which means
                                        # it was either deleted or modified. In either case,
                                        # this is a changed line.
                                        found_changed_line = True
                                    break

                            # Only add trailing comments if they're NOT immediately before a changed line
                            if not found_changed_line:
                                result_lines.extend(trailing_lines)
                    else:
                        # Keep it (shouldn't happen but be safe)
                        result_lines.append(opt_line)

        # Parse the cleaned code back into a CST module
        cleaned_code = "".join(result_lines)
        return cst.parse_module(cleaned_code)

    except Exception as e:
        logging.warning("Error cleaning comments: %s", e)
        sentry_sdk.capture_exception(e)
        return optimized_module


def clean_extraneous_comments_pipeline(
    original_module: cst.Module,
    optimized_code_and_explanations: list[CodeExplanationAndID],
    *,
    orig_code: str | None = None,
    orig_code_stripped: str | None = None,
) -> list[CodeExplanationAndID]:
    """Pipeline wrapper for comment cleaning.

    Cleans extraneous comments from all optimized code variants.
    Pre-computes original code and stripped version once for all candidates.
    """
    try:
        cleaned_results = []
        if orig_code is None:
            orig_code = original_module.code
        if orig_code_stripped is None:
            orig_code_stripped = _strip_comments_from_code(orig_code)

        for ce in optimized_code_and_explanations:
            try:
                cleaned_module = clean_extraneous_comments(
                    original_module, ce.cst_module, orig_code=orig_code, orig_code_stripped=orig_code_stripped
                )
                cleaned_results.append(
                    CodeExplanationAndID(cst_module=cleaned_module, explanation=ce.explanation, id=ce.id)
                )
            except Exception as e:
                logging.warning("Error cleaning comments for optimization %s: %s", ce.id, e)
                sentry_sdk.capture_exception(e)
                # Keep the original if cleaning fails
                cleaned_results.append(ce)

        return cleaned_results
    except Exception as e:
        logging.warning("Error in comment cleaning pipeline: %s", e)
        sentry_sdk.capture_exception(e)
        return optimized_code_and_explanations


def fix_forward_references(
    _original_module: cst.Module, optimized_code_and_explanations: list[CodeExplanationAndID]
) -> list[CodeExplanationAndID]:
    """Add 'from __future__ import annotations' to fix forward reference errors.

    This handles cases where type annotations reference names that aren't defined
    or imported in the code snippet. Without this import, Python would raise a
    NameError when trying to evaluate the annotation at runtime.
    """
    new_optimizations = []
    for ce in optimized_code_and_explanations:
        try:
            new_module = add_future_annotations_import(ce.cst_module)
            if new_module is not ce.cst_module:
                new_optimizations.append(
                    CodeExplanationAndID(cst_module=new_module, explanation=ce.explanation, id=ce.id)
                )
            else:
                new_optimizations.append(ce)
        except Exception as e:  # noqa: BLE001
            logging.warning("Error fixing forward references for optimization %s: %s", ce.id, e)
            sentry_sdk.capture_exception(e)
            # Keep the original if fixing fails
            new_optimizations.append(ce)
    return new_optimizations


def optimizations_postprocessing_pipeline(
    original_module: cst.Module, optimized_code_and_explanations: list[CodeExplanationAndID]
) -> list[CodeExplanationAndID]:
    # Pre-compute original code string once — avoids redundant CST codegen across steps
    original_code = original_module.code
    original_code_stripped = _strip_comments_from_code(original_code)

    optimized_code_and_explanations = fix_missing_docstring(original_module, optimized_code_and_explanations)
    optimized_code_and_explanations = clean_extraneous_comments_pipeline(
        original_module,
        optimized_code_and_explanations,
        orig_code=original_code,
        orig_code_stripped=original_code_stripped,
    )
    optimized_code_and_explanations = fix_forward_references(original_module, optimized_code_and_explanations)
    optimized_code_and_explanations = deduplicate_optimizations(original_module, optimized_code_and_explanations)
    optimized_code_and_explanations = equality_check(
        original_module, optimized_code_and_explanations, original_code=original_code
    )
    optimized_code_and_explanations = dedup_and_sort_imports(original_module, optimized_code_and_explanations)
    optimized_code_and_explanations = cleanup_explanations(original_module, optimized_code_and_explanations)
    optimized_code_and_explanations = filter_ellipsis_containing_code(original_module, optimized_code_and_explanations)
    return optimized_code_and_explanations
