from pathlib import Path
from typing import List, Optional, Literal, Dict, Tuple

from tree_sitter import Parser, Language, Node
import tree_sitter_python as tspython

from .schema import CodeChunk, ChunkAST, ChunkSpan, ChunkHierarchy, ChunkType

# ----------------------------
# Types
# ----------------------------

TS_TO_CHUNK_TYPE: Dict[str, ChunkType] = {
    "module": "module",
    "class_definition": "class",
    "function_definition": "function",
    "async_function_definition": "function",
    "import_statement": "imports",
    "import_from_statement": "imports",
}

MAX_TS_DEPTH = 3  # module → imports → class/function → method


# ----------------------------
# Helpers
# ----------------------------

def _safe_decode(data: bytes) -> str:
    try:
        return data.decode("utf-8")
    except UnicodeDecodeError:
        return data.decode("utf-8", errors="ignore")


def _get_node_name(node: Node) -> Optional[str]:
    """
    Extract identifier name for class / function nodes.
    """
    for child in node.children:
        if child.type == "identifier":
            text = child.text
            if isinstance(text, (bytes, bytearray)):
                return _safe_decode(text)
    return None


# ----------------------------
# Public API
# ----------------------------

def extract_ts_chunks(file_path: Path) -> List[CodeChunk]:
    source_bytes = file_path.read_bytes()

    language: Language = Language(tspython.language())
    parser = Parser()
    parser.language = language

    tree = parser.parse(source_bytes)
    root = tree.root_node

    chunks: List[CodeChunk] = []

    def walk(node: Node, depth: int = 0, parent_node: Optional[Node] = None) -> None:
        if depth > MAX_TS_DEPTH:
            return

        node_type = node.type

        if node_type in TS_TO_CHUNK_TYPE:
            code_bytes = source_bytes[node.start_byte : node.end_byte]
            code = _safe_decode(code_bytes)
            
            chunk_type = TS_TO_CHUNK_TYPE[node_type]
            name = _get_node_name(node)
            
            # For imports, use the full import as name
            if chunk_type == "imports":
                name = code.strip()
            
            # Create chunk with byte-level precision
            chunks.append(
                CodeChunk(
                    chunk_id=f"ts_{node.start_byte}_{node.end_byte}",
                    file_path=str(file_path),
                    language="python",
                    chunk_type=chunk_type,
                    code=code,
                    ast=ChunkAST(
                        symbol_type=None,  # TS doesn't provide semantic types
                        name=name,
                        parent=None,  # Parent relationships from AST
                        docstring=None,
                        decorators=[],
                        imports=[],
                        node_type=node_type,
                    ),
                    span=ChunkSpan(
                        start_byte=node.start_byte,
                        end_byte=node.end_byte,
                        start_line=node.start_point[0] + 1,
                        end_line=node.end_point[0] + 1,
                        char_count=len(code),
                    ),
                    hierarchy=ChunkHierarchy(
                        is_primary=False,  # Tree-sitter chunks are for span enrichment only
                        is_extracted=True,
                        depth=depth,
                        parent_id=None,  # Parent relationships from AST
                    ),
                    metadata={
                        "byte_span": {
                            "start": node.start_byte,
                            "end": node.end_byte,
                        },
                        "tree_sitter_node_type": node_type,
                        "is_exact_span": True,
                    },
                )
            )

        for child in node.children:
            walk(child, depth + 1, node)

    walk(root)
    return chunks