from pathlib import Path
import pandas as pd
import numpy as np
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter, MarkdownHeaderTextSplitter

SECTION_HEADERS_TO_SPLIT_ON = [
        ("##", "Header 2")
    ]
SECTION_MARKDOWN_SPLITTER = MarkdownHeaderTextSplitter(
        headers_to_split_on = SECTION_HEADERS_TO_SPLIT_ON
    )

SUBSECTION_HEADERS_TO_SPLIT_ON = [
        ("###", "Header 3"),
        ("####", "Header 4")
    ]
SUBSECTION_MARKDOWN_SPLITTER = MarkdownHeaderTextSplitter(
        headers_to_split_on = SUBSECTION_HEADERS_TO_SPLIT_ON
    )

MARKDOWN_SEPARATORS = [
    "\n#{1,6} ",
    "```\n",
    "\n\\*\\*\\*+\n",
    "\n---+\n",
    "\n___+\n",
    "\n\n",
    "\n",
    " ",
    "",
]

RECURSIVE_SPLITTER = RecursiveCharacterTextSplitter(
    chunk_size = 1000,
    chunk_overlap = 100,
    add_start_index = True,
    strip_whitespace = True,
    separators = MARKDOWN_SEPARATORS
)

MAX_CHARS = 1500

def read_markdown(file_name: str, parsed_dir:Path) -> str:
    """
    Reads a Markdown file from the parsed directory and returns its contents as a string.

    Args:
        file_name (str): Name of the Markdown file.

    Returns:
        str: The contents of the Markdown file.
    """
    markdown_path = parsed_dir / file_name

    with open(markdown_path, "r", encoding="utf-8") as f:
        return f.read()

def split_sections(markdown_text: str, game:str, source_file: str ) -> list[Document]:
    """
    Splits a Markdown document into parent sections based on level-two (##) headings.

    Args:
        markdown_text (str): The Markdown document to split.
        game (str): The name of game.
        source_file (str): The name of the source file.

    Returns:
        list[Document]: A list of LangChain Document objects, where each document
        represents a parent section and includes the associated heading metadata.
    """
    sections = SECTION_MARKDOWN_SPLITTER.split_text(markdown_text)
    game_id = game.lower()

    for i, section in enumerate(sections, start=1):
        section_name = section.metadata.get("Header 2", "Intro")
        section.metadata["game"] = game
        section.metadata["section_id"] = f"{game_id}_{i:03d}"
        section.metadata["section_name"] = section_name
        section.metadata["source"] = source_file

    return sections

def split_subsections(section: Document) -> list[Document]:
    """
    Splits a LangChain Document into subsections based on level-three (###) and level-four (####) headings.

    Args:
        section (Document): The LangChain Document representing a parent section.
    
    Returns:
        list[Document]: A list of LangChain Document objects, where each document
        represents a subsection and includes the associated heading metadata
    """
    subsections = SUBSECTION_MARKDOWN_SPLITTER.split_text(section.page_content)

    for doc in subsections:
        doc.metadata.update(section.metadata)
    
    for i, subsection in enumerate(subsections, start=1):
        subsection_name = subsection.metadata.get("Header 3", "General")
        topic_name = subsection.metadata.get("Header 4")
        subsection.metadata["subsection_id"] = f"SS{i:03d}"
        subsection.metadata["subsection_name"] = subsection_name
        subsection.metadata["topic_name"] = topic_name

    return subsections

def split_chunks(subsection: Document) -> list[Document]:
    """
    If the subsection is short enough, keep it as one chunk.
    If it is too long, split it into smaller chunks.
    """
    was_split = False

    if len(subsection.page_content) <= MAX_CHARS:
        chunks = [subsection]
    else:
        chunks = RECURSIVE_SPLITTER.split_documents([subsection])
        was_split = True
    
    for i, chunk in enumerate(chunks, start=1):
        chunk.metadata["chunk_id"] = (
            f'{chunk.metadata["section_id"]}_'
            f'{chunk.metadata["subsection_id"]}_'
            f'C{i:03d}'
        )
        chunk.metadata["chunk_order"] = i
        chunk.metadata["was_split"] = was_split

    return chunks