"""Python AST 分块"""

from __future__ import annotations

import ast

from .base import BaseChunker, Chunk, estimate_tokens
from .text import TextChunker

MAX_CLASS_TOKENS = 2000


class PythonChunker(BaseChunker):
    """用 ast 模块按函数/类分割 Python 文件"""

    def chunk(self, content: str, file_path: str) -> list[Chunk]:
        try:
            tree = ast.parse(content)
        except SyntaxError:
            return TextChunker().chunk(content, file_path)

        lines = content.splitlines(keepends=True)
        chunks: list[Chunk] = []

        # 收集顶层节点
        top_nodes: list[ast.AST] = []
        header_end = 0  # module 级代码的结束行

        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                top_nodes.append(node)
            elif hasattr(node, "end_lineno"):
                header_end = max(header_end, node.end_lineno)

        # Module header（import、常量等）
        if header_end > 0:
            header_text = "".join(lines[:header_end]).strip()
            if header_text:
                chunks.append(Chunk(
                    content=header_text,
                    section="module_header",
                    chunk_index=0,
                ))

        # 函数和类
        for node in top_nodes:
            start = node.lineno - 1
            end = node.end_lineno
            node_text = "".join(lines[start:end]).strip()

            if not node_text:
                continue

            section = f"class {node.name}" if isinstance(node, ast.ClassDef) else f"def {node.name}"

            # 超大 class 按方法分割
            if isinstance(node, ast.ClassDef) and estimate_tokens(node_text) > MAX_CLASS_TOKENS:
                chunks.extend(self._split_class(node, lines, len(chunks)))
            else:
                chunks.append(Chunk(
                    content=node_text,
                    section=section,
                    chunk_index=len(chunks),
                ))

        # 没有解析到任何节点时降级
        if not chunks:
            return TextChunker().chunk(content, file_path)

        return chunks

    def _split_class(self, class_node: ast.ClassDef, lines: list[str],
                     start_index: int) -> list[Chunk]:
        """将超大 class 按方法分割"""
        chunks: list[Chunk] = []
        class_name = f"class {class_node.name}"

        # class 定义行 + docstring
        first_method = None
        for node in ast.iter_child_nodes(class_node):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                first_method = node
                break

        if first_method:
            header_end = first_method.lineno - 1
            header_text = "".join(lines[class_node.lineno - 1:header_end]).strip()
            if header_text:
                chunks.append(Chunk(
                    content=header_text,
                    section=f"{class_name} (header)",
                    chunk_index=start_index + len(chunks),
                ))

        # 每个方法一个 chunk
        for node in ast.iter_child_nodes(class_node):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                method_text = "".join(lines[node.lineno - 1:node.end_lineno]).strip()
                if method_text:
                    chunks.append(Chunk(
                        content=method_text,
                        section=f"{class_name}.{node.name}",
                        parent_section=class_name,
                        chunk_index=start_index + len(chunks),
                    ))

        return chunks
