from __future__ import annotations
from typing import List, Dict, Any, Optional, TypedDict, Tuple
from markdown_it import MarkdownIt
from markdown_it.tree import SyntaxTreeNode
import traceback
import logging
import json

import defines

class Chunk(TypedDict):
    text: str
    metadata: Dict[str, Any]

def clear_chunk(chunk: Chunk):
    chunk["text"] = ""
    chunk["metadata"] = {
        "doc_type": "unknown",
        "source_file": chunk["metadata"]["source_file"],
        "lines": chunk["metadata"]["lines"],
        "path": "",  # This will be updated during processing
        "level": 0,
    }
    return chunk

class MarkdownChunker:
    def __init__(self):
        # Initialize the Markdown parser
        self.md_parser = MarkdownIt("commonmark")

    def process_file(self, file_path: str) -> Optional[List[Chunk]]:
        """
        Process a single markdown file and return chunks.

        Args:
            file_path: Path to the markdown file

        Returns:
            List of chunks with metadata or None if file can't be processed
        """
        try:
            # Read the markdown file
            with open(file_path, "r", encoding="utf-8") as f:
                content = f.read()

            # Parse the markdown
            tokens = self.md_parser.parse(content)
            logging.info(f"Found {len(tokens)} in {file_path}")

            ast = SyntaxTreeNode(tokens)

            # Extract chunks with metadata
            chunks = self.extract_chunks(ast, file_path, len(content.splitlines()))

            return chunks

        except Exception as e:
            logging.error(f"Error processing {file_path}: {str(e)}")
            logging.error(traceback.format_exc())

            return None

    def extract_chunks(self, ast: SyntaxTreeNode, file_path: str, total_lines: int) -> List[Chunk]:
        """
        Extract logical chunks from the AST with appropriate metadata.

        Args:
            ast: The Abstract Syntax Tree from markdown-it-py
            file_path: Path to the source file

        Returns:
            List of chunks with metadata
        """
        chunks: List[Chunk] = []
        current_headings: List[str] = [""] * 6  # Track h1-h6 headings

        # Initialize a chunk structure
        chunk: Chunk = {
            "text": "",
            "metadata": {
                "source_file": file_path,
                "lines": total_lines
            },
        }
        clear_chunk(chunk)

        # Process the AST recursively
        self._process_node(ast, current_headings, chunks, chunk, level=0)

        return chunks

    def _sanitize_metadata(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
        return {
            k: ("" if v is None else v) for k, v in metadata.items() if v is not None
        }

    def _extract_text_from_children(self, node: SyntaxTreeNode) -> str:
        lines = []
        for child in node.children:
            if child.type == "list_item":
                lines.append(f"- {self._extract_text_from_children(child).strip()}")
            elif child.type == "fence":
                info = f"{child.info}" if hasattr(child, "info") and child.info else ""
                lines.append(f"\n```{info}\n{child.content.strip()}\n```\n")
            elif child.children:
                lines.append(self._extract_text_from_children(child).strip())
            elif hasattr(child, "content"):
                lines.append(child.content.strip())
        return "\n".join(lines)

    def _process_node(
        self,
        node: SyntaxTreeNode,
        current_headings: list[str],
        chunks: List[Chunk],
        chunk: Chunk,
        level: int,
        buffer: int = defines.chunk_buffer
    ) -> int:
        is_list = False
        # Handle heading nodes
        if node.type == "heading":
            # print(f'{" " * level}{chunk["metadata"]["path"]} {node.type}')
            if node.is_nested and node.nester_tokens:
                opening, closing = node.nester_tokens
                level = int(opening.tag[1:]) - 1

            heading_text = self._extract_text_from_children(node)

            # Update lower heading states
            current_headings[level] = heading_text
            for i in range(level, 6):
                if i != level:
                    current_headings[i] = ""

            # Store previous chunk if it has text
            if chunk["text"].strip():
                path = " > ".join([h for h in current_headings if h])
                chunk["text"] = chunk["text"].strip()
                chunk["metadata"]["path"] = path
                chunk["metadata"]["level"] = level
                if node.nester_tokens:
                    opening, closing = node.nester_tokens
                    if opening and opening.map:
                        ( begin, end ) = opening.map
                        metadata = chunk["metadata"]
                        metadata["chunk_begin"] = max(0, begin - buffer)
                        metadata["chunk_end"] = min(metadata["lines"], end + buffer)
                        metadata["line_begin"] = begin
                        metadata["line_end"] = end

                chunks.append(chunk.copy())
                clear_chunk(chunk)

        # Add code block directly to current chunk
        elif node.type == "fence":
            # print(f'{" " * level}{chunk["metadata"]["path"]} {node.type}')
            language = node.info.strip() if hasattr(node, "info") and node.info else ""
            code_block = f"\n```{language}\n{node.content.strip()}\n```\n"
            chunk["text"] += code_block

        # Handle list structures
        # elif node.type in ["list", "list_item"]:
        #     # print(node.nester_tokens)
        #     # print(node.pretty(show_text=True))
        #     is_list = True
        #     list_chunk = []
        #     for child in node.children:
        #         level = self._process_node(child, current_headings, chunks, chunk, level=level)

        # Handle paragraph
        elif node.type in ["paragraph"]:  # , "list", "list_item"]:
            text = self._extract_text_from_children(node)
            # indented_text = "\n".join([f'{"-" * (level + 2)}{line}' for line in text.split('\n')])
            # print(f'{"-" * level}{chunk["metadata"]["path"]} {node.type}:\n{indented_text} {len(node.children)}\n')
            text = text.strip()
            if text:
                # indented_text = "\n".join([f'{"-" * (level + 2)}{line}' for line in text.split('\n')])
                # print(f'{" " * level}{chunk["metadata"]["path"]} {node.type}:\n{indented_text}\n')
                chunk["text"] += f"\n{text}\n"
                chunk["text"] = chunk["text"].strip()
                if chunk["text"]:
                    path = " > ".join([h for h in current_headings if h])
                    chunk["text"] = chunk["text"]
                    chunk["metadata"]["path"] = path
                    chunk["metadata"]["level"] = level
                    if node.nester_tokens:
                        opening, closing = node.nester_tokens
                        if opening and opening.map:
                            ( begin, end ) = opening.map
                            metadata = chunk["metadata"]
                            metadata["chunk_begin"] = max(0, begin - buffer)
                            metadata["chunk_end"] = min(metadata["lines"], end + buffer)
                            metadata["line_begin"] = begin
                            metadata["line_end"] = end
                    chunks.append(chunk.copy())
                    clear_chunk(chunk)

        # Recursively process children
        if not is_list:
            for child in node.children:
                level = self._process_node(
                    child, current_headings, chunks, chunk, level=level
                )

        # After root-level recursion, finalize any remaining chunk
        if node.type == "document":
            # print(node.type)
            if chunk["text"].strip():
                path = " > ".join([h for h in current_headings if h])
                chunk["metadata"]["path"] = path
                if node.nester_tokens:
                    opening, closing = node.nester_tokens
                    if opening and opening.map:
                        ( begin, end ) = opening.map
                        metadata = chunk["metadata"]
                        metadata["chunk_begin"] = max(0, begin - buffer)
                        metadata["chunk_end"] = min(metadata["lines"], end + buffer)
                        metadata["line_begin"] = begin
                        metadata["line_end"] = end
                chunks.append(chunk.copy())

        return level
