from __future__ import annotations

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

import attrs
import isort
import libcst as cst
from libcst import (
    BaseStatement,
    CSTTransformer,
    CSTVisitor,
    Expr,
    IndentedBlock,
    SimpleStatementLine,
    SimpleString,
)
from libcst.codemod import CodemodContext
from libcst.codemod.visitors import AddImportsVisitor

from codeflash_api.languages.python._cst_utils import (
    collect_imported_names_from_import,
    collect_imported_names_from_import_from,
    compare_unparsed_ast_to_source,
    get_dotted_name,
    parse_module_to_cst,
    unparse_parse_source,
)

if TYPE_CHECKING:
    from libcst import FunctionDef

log = logging.getLogger(__name__)

BUILTIN_NAMES = frozenset(dir(builtins))


@attrs.frozen
class OptimizationCandidate:
    """
    An optimized code variant with its CST module and explanation.
    """

    cst_module: cst.Module
    explanation: str
    id: str


def safe_isort(code: str, **kwargs: Any) -> str:
    """
    Run isort on *code*, returning original on failure.
    """
    try:
        return isort.code(code, **kwargs)
    except Exception:
        return code


def deduplicate_optimizations(
    _original_module: cst.Module,
    candidates: list[OptimizationCandidate],
) -> list[OptimizationCandidate]:
    """
    Remove candidates with equivalent ASTs.
    """
    seen_asts: set[str] = set()
    unique: list[OptimizationCandidate] = []
    for candidate in candidates:
        normalized = ast.unparse(ast.parse(candidate.cst_module.code))
        if normalized not in seen_asts:
            seen_asts.add(normalized)
            unique.append(candidate)
    return unique


def equality_check(
    original_module: cst.Module,
    candidates: list[OptimizationCandidate],
    *,
    original_code: str | None = None,
) -> list[OptimizationCandidate]:
    """
    Filter out candidates identical to the original.
    """
    source = (
        original_code if original_code is not None else original_module.code
    )
    try:
        original_ast = unparse_parse_source(source)
    except Exception:
        return [c for c in candidates if c.cst_module.code != source]
    filtered: list[OptimizationCandidate] = []
    for c in candidates:
        try:
            if not compare_unparsed_ast_to_source(
                original_ast, c.cst_module.code
            ):
                filtered.append(c)
        except Exception:
            if c.cst_module.code != source:
                filtered.append(c)
    return filtered


_EXPLANATION_PATTERNS = [
    (
        re.compile(
            r"\sHere (is|are) (the )?((code|optimization)"
            r"|\S+\scode|(optimized|improved) versions?"
            r" 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,
    candidates: list[OptimizationCandidate],
) -> list[OptimizationCandidate]:
    """
    Clean up explanation text from LLM artifacts.
    """
    result: list[OptimizationCandidate] = []
    for c in candidates:
        cleaned = c.explanation
        for pattern, repl in _EXPLANATION_PATTERNS:
            cleaned = pattern.sub(repl, cleaned)
        result.append(
            OptimizationCandidate(
                cst_module=c.cst_module,
                explanation=cleaned,
                id=c.id,
            )
        )
    return result


class DocstringVisitor(CSTVisitor):
    """
    Collect original docstrings from functions and classes.
    """

    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:
        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):
    """
    Restore original docstrings in optimized code.
    """

    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:
        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:
        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,
    candidates: list[OptimizationCandidate],
) -> list[OptimizationCandidate]:
    """
    Restore docstrings that the LLM removed.
    """
    visitor = DocstringVisitor()
    try:
        original_module.visit(visitor)
    except Exception:
        return candidates
    transformer = DocstringTransformer(visitor.original_docstrings)
    result: list[OptimizationCandidate] = []
    for c in candidates:
        try:
            result.append(
                OptimizationCandidate(
                    cst_module=c.cst_module.visit(transformer),
                    explanation=c.explanation,
                    id=c.id,
                )
            )
        except Exception:
            log.warning("Failed to restore docstring for %s", c.id)
            result.append(c)
    return result


def dedup_and_sort_imports(
    _original_module: cst.Module,
    candidates: list[OptimizationCandidate],
) -> list[OptimizationCandidate]:
    """
    Run isort on all candidates.
    """
    result: list[OptimizationCandidate] = []
    for c in candidates:
        try:
            original_code = c.cst_module.code
            sorted_code = safe_isort(original_code, disregard_skip=True)
        except Exception:
            result.append(c)
            continue
        if sorted_code == original_code:
            result.append(c)
        else:
            result.append(
                OptimizationCandidate(
                    cst_module=parse_module_to_cst(sorted_code),
                    explanation=c.explanation,
                    id=c.id,
                )
            )
    return result


class EllipsisContainingCodeVisitor(CSTVisitor):
    """
    Detect any ellipsis in code.
    """

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

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


def filter_ellipsis_containing_code(
    original_module: cst.Module,
    candidates: list[OptimizationCandidate],
) -> list[OptimizationCandidate]:
    """
    Filter out candidates that introduce ellipsis not present in original.
    """
    original_visitor = EllipsisContainingCodeVisitor()
    original_module.visit(original_visitor)
    if original_visitor.ellipsis_containing_code:
        return candidates
    result: list[OptimizationCandidate] = []
    for c in candidates:
        visitor = EllipsisContainingCodeVisitor()
        c.cst_module.visit(visitor)
        if not visitor.ellipsis_containing_code:
            result.append(c)
    return result


def _strip_comments_from_code(code: str) -> str:
    """
    Remove all comments from Python code, preserving strings.
    """
    try:
        lines = code.splitlines(keepends=True)
        tokens = tokenize.generate_tokens(io.StringIO(code).readline)

        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])
                    )

        result_lines: list[str] = []
        for line_num, line in enumerate(lines, start=1):
            current = line
            line_comments = comments_by_line[line_num - 1]

            if line_comments:
                line_comments.sort(reverse=True)
                for start_col, end_col in line_comments:
                    current = current[:start_col].rstrip() + current[end_col:]

                if not current.endswith("\n") and lines[line_num - 1].endswith(
                    "\n"
                ):
                    current += "\n"

            result_lines.append(current)

        return "".join(result_lines)
    except (tokenize.TokenError, IndentationError):
        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:
    """
    Remove extraneous comments from optimized code using diff-based approach.
    """
    try:
        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)

        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)

        orig_code_only = orig_code_stripped.splitlines(keepends=True)
        opt_code_only = opt_code_stripped.splitlines(keepends=True)

        orig_code_line_indices: list[int] = []
        orig_code_lines_filtered: list[str] = []
        for i, line in enumerate(orig_code_only):
            if line.strip():
                orig_code_line_indices.append(i)
                orig_code_lines_filtered.append(line)

        opt_code_line_indices: list[int] = []
        opt_code_lines_filtered: list[str] = []
        for i, line in enumerate(opt_code_only):
            if line.strip():
                opt_code_line_indices.append(i)
                opt_code_lines_filtered.append(line)

        code_matcher = SequenceMatcher(
            None,
            orig_code_lines_filtered,
            opt_code_lines_filtered,
        )
        code_changed_line_indices: set[int] = set()

        for tag, _i1, _i2, j1, j2 in code_matcher.get_opcodes():
            if tag != "equal":
                for j in range(j1, j2):
                    code_changed_line_indices.add(j)

        code_changed_lines: set[int] = 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])

        orig_to_opt_mapping: dict[int, int] = {}

        for tag, i1, i2, j1, j2 in code_matcher.get_opcodes():
            if tag == "equal":
                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

        result_lines: list[str] = []
        orig_idx = 0
        restored_orig_indices: set[int] = set()

        for opt_idx, opt_line in enumerate(opt_lines):
            if opt_idx in code_changed_lines:
                result_lines.append(opt_line)
            else:
                opt_code = (
                    opt_code_only[opt_idx]
                    if opt_idx < len(opt_code_only)
                    else ""
                )

                is_comment_only = not opt_code.strip()
                is_near_change = False
                if is_comment_only:
                    for check_idx in range(opt_idx + 1, len(opt_lines)):
                        if check_idx in code_changed_lines:
                            is_near_change = True
                            break
                        check_code = (
                            opt_code_only[check_idx]
                            if check_idx < len(opt_code_only)
                            else ""
                        )
                        if check_code.strip():
                            break

                    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
                            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:
                    result_lines.append(opt_line)
                elif is_comment_only and not is_near_change:
                    pass
                elif opt_idx in code_changed_lines:
                    result_lines.append(opt_line)
                else:
                    found_orig = None
                    orig_line_idx = None
                    for orig_idx_search in range(orig_idx, len(orig_lines)):
                        orig_code_line = (
                            orig_code_only[orig_idx_search]
                            if orig_idx_search < len(orig_code_only)
                            else ""
                        )
                        if orig_code_line == 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:
                        if (
                            orig_line_idx is not None
                            and orig_line_idx > 0
                            and opt_idx not in code_changed_lines
                        ):
                            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():
                                    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:
                                    break

                            result_lines.extend(preceding_lines)

                        result_lines.append(found_orig)
                        if orig_line_idx is not None:
                            restored_orig_indices.add(orig_line_idx)

                        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: list[str] = []
                            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():
                                    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:
                                    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:
                                        found_changed_line = True
                                    break

                            if not found_changed_line:
                                result_lines.extend(trailing_lines)
                    else:
                        result_lines.append(opt_line)

        cleaned_code = "".join(result_lines)
        return cst.parse_module(cleaned_code)

    except Exception:
        log.warning("Error cleaning comments")
        return optimized_module


def clean_extraneous_comments_pipeline(
    original_module: cst.Module,
    candidates: list[OptimizationCandidate],
    *,
    orig_code: str | None = None,
    orig_code_stripped: str | None = None,
) -> list[OptimizationCandidate]:
    """
    Pipeline wrapper for comment cleaning.
    """
    try:
        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)

        result: list[OptimizationCandidate] = []
        for c in candidates:
            try:
                cleaned = clean_extraneous_comments(
                    original_module,
                    c.cst_module,
                    orig_code=orig_code,
                    orig_code_stripped=orig_code_stripped,
                )
                result.append(
                    OptimizationCandidate(
                        cst_module=cleaned,
                        explanation=c.explanation,
                        id=c.id,
                    )
                )
            except Exception:
                log.warning(
                    "Error cleaning comments for %s",
                    c.id,
                )
                result.append(c)
        return result
    except Exception:
        log.warning("Error in comment cleaning pipeline")
        return candidates


def extract_names_from_cst_annotation(
    annotation: cst.BaseExpression,
    names: set[str],
) -> None:
    """
    Walk a CST annotation expression and collect Name nodes.
    """
    if isinstance(annotation, cst.Name):
        names.add(annotation.value)
    elif isinstance(annotation, cst.Attribute):
        extract_names_from_cst_annotation(annotation.value, names)
    elif isinstance(annotation, cst.Subscript):
        extract_names_from_cst_annotation(annotation.value, names)
        for element in annotation.slice:
            if isinstance(element, cst.SubscriptElement):
                slice_node = element.slice
                if isinstance(slice_node, cst.Index):
                    extract_names_from_cst_annotation(slice_node.value, names)
    elif isinstance(annotation, cst.BinaryOperation):
        extract_names_from_cst_annotation(annotation.left, names)
        extract_names_from_cst_annotation(annotation.right, names)
    elif isinstance(annotation, (cst.Tuple, cst.List)):
        for el in annotation.elements:
            if isinstance(
                el,
                (cst.Element, cst.StarredElement),
            ):
                extract_names_from_cst_annotation(el.value, names)


class CSTAnnotationNameCollector(cst.CSTVisitor):
    """
    Collect names used in type annotations.
    """

    def __init__(self) -> None:
        self.annotation_names: set[str] = set()
        self.defined_names: set[str] = set()
        self.imported_names: set[str] = set()

    def visit_Import(self, node: cst.Import) -> bool | None:
        self.imported_names.update(collect_imported_names_from_import(node))
        return False

    def visit_ImportFrom(self, node: cst.ImportFrom) -> bool | None:
        self.imported_names.update(
            collect_imported_names_from_import_from(node)
        )
        return False

    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool | None:
        self.defined_names.add(node.name.value)
        self._collect_from_function(node)
        return True

    def visit_ClassDef(self, node: cst.ClassDef) -> bool | None:
        self.defined_names.add(node.name.value)
        return True

    def visit_AnnAssign(self, node: cst.AnnAssign) -> bool | None:
        extract_names_from_cst_annotation(
            node.annotation.annotation,
            self.annotation_names,
        )
        return True

    def _collect_from_function(self, node: cst.FunctionDef) -> None:
        """
        Collect annotation names from function parameters and return type.
        """
        params = node.params
        for param in [
            *params.params,
            *params.kwonly_params,
            *params.posonly_params,
        ]:
            if param.annotation:
                extract_names_from_cst_annotation(
                    param.annotation.annotation,
                    self.annotation_names,
                )
        if (
            params.star_arg
            and isinstance(params.star_arg, cst.Param)
            and params.star_arg.annotation
        ):
            extract_names_from_cst_annotation(
                params.star_arg.annotation.annotation,
                self.annotation_names,
            )
        if params.star_kwarg and params.star_kwarg.annotation:
            extract_names_from_cst_annotation(
                params.star_kwarg.annotation.annotation,
                self.annotation_names,
            )
        if node.returns:
            extract_names_from_cst_annotation(
                node.returns.annotation,
                self.annotation_names,
            )

    def get_undefined_annotation_names(self) -> set[str]:
        """
        Return annotation names not defined, imported, or builtin.
        """
        return (
            self.annotation_names
            - self.defined_names
            - self.imported_names
            - BUILTIN_NAMES
        )


def has_future_annotations_import(
    module: cst.Module,
) -> bool:
    """
    Check if *module* has 'from __future__ import annotations'.
    """
    for stmt in module.body:
        if isinstance(stmt, cst.SimpleStatementLine):
            for small_stmt in stmt.body:
                if (
                    isinstance(small_stmt, cst.ImportFrom)
                    and small_stmt.module
                    and get_dotted_name(small_stmt.module) == "__future__"
                ):
                    if isinstance(small_stmt.names, cst.ImportStar):
                        continue
                    for alias in small_stmt.names:
                        if alias.name.value == "annotations":
                            return True
    return False


def add_future_annotations_import(
    module: cst.Module,
) -> cst.Module:
    """
    Add 'from __future__ import annotations' if needed for forward refs.
    """
    if has_future_annotations_import(module):
        return module

    collector = CSTAnnotationNameCollector()
    module.visit(collector)
    if not collector.get_undefined_annotation_names():
        return module

    context = CodemodContext()
    AddImportsVisitor.add_needed_import(context, "__future__", "annotations")
    return module.visit(AddImportsVisitor(context))


def fix_forward_references(
    _original_module: cst.Module,
    candidates: list[OptimizationCandidate],
) -> list[OptimizationCandidate]:
    """
    Add future annotations import where needed.
    """
    result: list[OptimizationCandidate] = []
    for c in candidates:
        try:
            new_module = add_future_annotations_import(c.cst_module)
            if new_module is not c.cst_module:
                result.append(
                    OptimizationCandidate(
                        cst_module=new_module,
                        explanation=c.explanation,
                        id=c.id,
                    )
                )
            else:
                result.append(c)
        except Exception:
            log.warning("Error fixing forward refs for %s", c.id)
            result.append(c)
    return result


def optimizations_postprocessing_pipeline(
    original_module: cst.Module,
    candidates: list[OptimizationCandidate],
) -> list[OptimizationCandidate]:
    """
    Run the full postprocessing pipeline on optimization candidates.
    """
    original_code = original_module.code
    original_code_stripped = _strip_comments_from_code(original_code)

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