from typing import List, Dict, Any, Optional
import logging
import uuid
from langchain_text_splitters import (
    MarkdownHeaderTextSplitter,
    RecursiveCharacterTextSplitter,
)

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


IRRELEVANT_HEADERS = [
    "reference",
    "references",
    "bibliography",
    "daftar pustaka",
    "referensi",
    "acknowledgement",
    "acknowledgements",
    "acknowledgment",
    "acknowledgments",
    "ucapan terima kasih",
    "funding",
    "pendanaan",
    "appendix",
    "appendices",
    "lampiran",
    "table of contents",
    "daftar isi",
    "author contributions",
    "kontribusi penulis"
]


def _normalize_header(h: Optional[str]) -> str:
    if not h:
        return ""
    return h.strip().lower()


def is_irrelevant_header(header: Optional[str]) -> bool:
    """Return True when the chunk header matches the filtering blacklist."""
    normalized = _normalize_header(header)
    if not normalized:
        return False
    return any(token in normalized for token in IRRELEVANT_HEADERS)


def build_llm_context_from_chunks(
    chunks: List[Dict[str, Any]],
    max_chars: int = 25000,
) -> tuple[str, List[Dict[str, Any]]]:
    """Filter out irrelevant chunks and join the remaining text for LLM input.

    Returns:
        tuple[str, List[Dict[str, Any]]]:
            - joined text bounded by max_chars
            - list of chunks that survived filtering
    """
    if max_chars <= 0:
        raise ValueError("max_chars must be greater than 0")

    filtered_chunks: List[Dict[str, Any]] = []
    text_parts: List[str] = []
    current_length = 0

    for chunk in chunks:
        header = chunk.get("metadata_header", "")
        if is_irrelevant_header(header):
            continue

        content = str(chunk.get("content", "") or "").strip()
        if not content:
            continue

        separator_length = 2 if text_parts else 0
        next_length = current_length + separator_length + len(content)
        if next_length > max_chars:
            remaining = max_chars - current_length - separator_length
            if remaining <= 0:
                break
            content = content[:remaining].rstrip()
            next_length = current_length + separator_length + len(content)

        filtered_chunks.append(chunk)
        text_parts.append(content)
        current_length = next_length

        if current_length >= max_chars:
            break

    return "\n\n".join(text_parts), filtered_chunks


def generate_chunks(
    markdown_text: str,
    file_hash: Optional[str] = None,
    title: Optional[str] = None,
    chunk_size: int = 1500,
    chunk_overlap: int = 150,
) -> List[Dict[str, Any]]:
    """
    Generate a list of chunk dictionaries from a markdown string.

    Each returned dict contains: `file_hash`, `title`, `chunk_index`,
    `metadata_header` and `content`.

    Algorithm:
    1. Use `MarkdownHeaderTextSplitter` to split markdown into header-aware blocks.
    2. For each block, apply `RecursiveCharacterTextSplitter` to enforce size & overlap.
    3. Preserve the nearest header as `metadata_header` for each chunk.

    Returns a list ordered by original document order.
    """
    if not isinstance(markdown_text, str):
        raise TypeError("markdown_text must be a string")

    results: List[Dict[str, Any]] = []

    # Step 1: header-aware split
    # Provide common header levels to the splitter to ensure deterministic behavior
    header_levels = [("#", "h1"), ("##", "h2"), ("###", "h3"), ("####", "h4"), ("#####", "h5")]
    header_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=header_levels)
    try:
        header_docs = header_splitter.split_text(markdown_text)
    except TypeError:
        # Fallback for versions that expose split_documents API
        header_docs = header_splitter.split_documents([{"text": markdown_text}])

    chunker = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size, chunk_overlap=chunk_overlap, add_start_index=True
    )

    chunk_index = 0

    for block in header_docs:
        # MarkdownHeaderTextSplitter returns Document objects with page_content and metadata.
        # Keep a compatibility fallback for alternate return shapes.
        if hasattr(block, "page_content"):
            block_text = str(getattr(block, "page_content", "") or "")
            metadata = getattr(block, "metadata", {}) or {}
            header_parts = []
            for key in ("h5", "h4", "h3", "h2", "h1"):
                value = metadata.get(key)
                if value:
                    header_parts.append(str(value).strip())
            header = " / ".join(header_parts)
        elif isinstance(block, dict):
            block_text = block.get("text", "")
            header = block.get("header", None) or block.get("metadata", {}).get("header")
        else:
            block_text = str(block)
            header = None

        normalized_header = _normalize_header(header)

        # Step 2: recursive split of block_text
        splits = chunker.split_text(block_text)

        for split in splits:
            # skip empty chunks
            content = split.strip() if isinstance(split, str) else str(split)
            if not content:
                continue

            results.append(
                {
                    "file_hash": file_hash,
                    "title": title,
                    "chunk_index": chunk_index,
                    "metadata_header": header if header else "",
                    "content": content,
                }
            )
            chunk_index += 1

    logger.info(f"Generated {len(results)} chunks (chunk_size={chunk_size}, overlap={chunk_overlap})")
    return results


def generate_parent_child_chunks(
    markdown_text: str,
    file_hash: Optional[str] = None,
    title: Optional[str] = None,
    parent_chunk_size: int = 1500,
    parent_chunk_overlap: int = 150,
    child_chunk_size: int = 200,
    child_chunk_overlap: int = 20,
) -> tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    """
    Generate two sets of chunks: large parent chunks for LLM context, and small child chunks for semantic search.
    Each child chunk contains a 'parent_id' linking it to its corresponding parent chunk.
    """
    if not isinstance(markdown_text, str):
        raise TypeError("markdown_text must be a string")

    parent_chunks: List[Dict[str, Any]] = []
    child_chunks: List[Dict[str, Any]] = []

    # Step 1: header-aware split
    header_levels = [("#", "h1"), ("##", "h2"), ("###", "h3"), ("####", "h4"), ("#####", "h5")]
    header_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=header_levels)
    try:
        header_docs = header_splitter.split_text(markdown_text)
    except TypeError:
        # Fallback for versions that expose split_documents API
        header_docs = header_splitter.split_documents([{"text": markdown_text}])

    parent_chunker = RecursiveCharacterTextSplitter(
        chunk_size=parent_chunk_size, chunk_overlap=parent_chunk_overlap, add_start_index=True
    )
    
    child_chunker = RecursiveCharacterTextSplitter(
        chunk_size=child_chunk_size, chunk_overlap=child_chunk_overlap, add_start_index=True
    )

    parent_index = 0
    child_index = 0

    for block in header_docs:
        if hasattr(block, "page_content"):
            block_text = str(getattr(block, "page_content", "") or "")
            metadata = getattr(block, "metadata", {}) or {}
            header_parts = []
            for key in ("h5", "h4", "h3", "h2", "h1"):
                value = metadata.get(key)
                if value:
                    header_parts.append(str(value).strip())
            header = " / ".join(header_parts)
        elif isinstance(block, dict):
            block_text = block.get("text", "")
            header = block.get("header", None) or block.get("metadata", {}).get("header")
        else:
            block_text = str(block)
            header = None

        normalized_header = _normalize_header(header)

        # Step 2: recursive split of block_text into parent chunks
        p_splits = parent_chunker.split_text(block_text)

        for p_split in p_splits:
            p_content = p_split.strip() if isinstance(p_split, str) else str(p_split)
            if not p_content:
                continue
                
            parent_id = str(uuid.uuid4())
            
            parent_chunks.append({
                "parent_id": parent_id,
                "file_hash": file_hash,
                "title": title,
                "chunk_index": parent_index,
                "metadata_header": header if header else "",
                "content": p_content,
            })
            parent_index += 1
            
            # Step 3: recursive split of parent chunk into child chunks
            c_splits = child_chunker.split_text(p_content)
            for c_split in c_splits:
                c_content = c_split.strip() if isinstance(c_split, str) else str(c_split)
                if not c_content:
                    continue
                    
                child_chunks.append({
                    "parent_id": parent_id,
                    "file_hash": file_hash,
                    "title": title,
                    "chunk_index": child_index,
                    "metadata_header": header if header else "",
                    "content": c_content,
                })
                child_index += 1

    logger.info(f"Generated {len(parent_chunks)} parent chunks and {len(child_chunks)} child chunks")
    return parent_chunks, child_chunks


if __name__ == "__main__":
    # Simple smoke test when run directly
    sample_md = """
# Title

## Introduction

This is a short introduction that will be split into chunks for testing. We will test the multi-level chunking functionality as well. """

    parents, children = generate_parent_child_chunks(sample_md, file_hash="testhash", title="Sample")
    print(f"Parents: {len(parents)}")
    for p in parents:
        print(f"  [Parent {p['chunk_index']}] ID: {p['parent_id']} - len: {len(p['content'])}")
        children_of_p = [c for c in children if c['parent_id'] == p['parent_id']]
        print(f"    Children: {len(children_of_p)}")
        for c in children_of_p:
            print(f"      [Child {c['chunk_index']}] len: {len(c['content'])}")
