"""Markdown chef for Chonkie."""

import os
import re
from pathlib import Path
from typing import Union

from chonkie.logger import get_logger
from chonkie.pipeline import chef
from chonkie.tokenizer import AutoTokenizer, TokenizerProtocol
from chonkie.types import (
    Chunk,
    MarkdownCode,
    MarkdownDocument,
    MarkdownImage,
    MarkdownTable,
)

from .base import BaseChef

logger = get_logger(__name__)


@chef("markdown")
class MarkdownChef(BaseChef):
    """Chef to process a markdown file into a MarkdownDocument type.

    Args:
      path: The path to the markdown file.

    Returns:
      MarkdownDocument: The processed markdown document.

    """

    def __init__(self, tokenizer: Union[TokenizerProtocol, str] = "character") -> None:
        """Initialize the MarkdownChef."""
        super().__init__()
        self.tokenizer = AutoTokenizer(tokenizer)
        self.code_pattern = re.compile(r"```([a-zA-Z0-9+\-_]*)\n?(.*?)\n?```", re.DOTALL)
        self.table_pattern = re.compile(r"(\|.*?\n\|[-: ]+\|.*?\n(?:\|.*?\n)*)")
        self.image_pattern = re.compile(r"(\[)?!\[([^\]]*)\]\(([^)]+)\)(?(1)\]\(([^)]+)\)|)")

    def prepare_tables(self, markdown: str) -> list[MarkdownTable]:
        """Prepare the tables for the MarkdownDocument.

        Args:
            markdown (str): The markdown text containing tables.

        Returns:
            list[MarkdownTable]: The list of tables with their start and end indices.

        """
        markdown_tables: list[MarkdownTable] = []
        for match in self.table_pattern.finditer(markdown):
            table_content = match.group(0)
            start_index = match.start()
            end_index = match.end()
            markdown_tables.append(
                MarkdownTable(content=table_content, start_index=start_index, end_index=end_index),
            )
        return markdown_tables

    def prepare_code(self, markdown: str) -> list[MarkdownCode]:
        """Extract markdown code snippets from a markdown string.

        Args:
            markdown (str): The markdown text containing code snippets.

        Returns:
            list[MarkdownCode]: A list of MarkdownCode objects, each containing
            the code content, language (if specified), and position indices.

        """
        # Pattern to capture language and content separately
        code_snippets: list[MarkdownCode] = []
        for match in self.code_pattern.finditer(markdown):
            language = match.group(1) if match.group(1) else None
            content = match.group(2)

            start_index = match.start()
            end_index = match.end()

            code_snippets.append(
                MarkdownCode(
                    content=content,
                    language=language,
                    start_index=start_index,
                    end_index=end_index,
                ),
            )
        return code_snippets

    def extract_images(self, markdown: str) -> list[MarkdownImage]:
        """Extract images from a markdown string.

        Args:
            markdown (str): The markdown text containing images.

        Returns:
            dict[str, str]: A dictionary where keys are image names (alt text or filename)
            and values are image paths or base64 data URLs.

        """
        images: list[MarkdownImage] = []

        for match in self.image_pattern.finditer(markdown):
            # Extract the match groups
            _, alt_text, image_src, link_url = match.groups()

            # Determine the key for the image
            if alt_text:
                key = alt_text
            else:
                # If no alt text, use filename from path
                if image_src.startswith("data:"):
                    # For base64 data URLs, use a generic name or extract from data URL
                    key = "base64_image"
                else:
                    # Extract filename from path
                    key = Path(image_src).name

            # Handle duplicate keys by appending a counter
            original_key = key
            counter = 1
            while key in images:
                key = f"{original_key}_{counter}"
                counter += 1

            images.append(
                MarkdownImage(
                    alias=key,
                    content=image_src,
                    start_index=match.start(),
                    end_index=match.end(),
                    link=link_url,
                ),
            )

        return images

    def extract_chunks(
        self,
        markdown: str,
        tables: list[MarkdownTable],
        code: list[MarkdownCode],
        images: list[MarkdownImage],
    ) -> list[Chunk]:
        """Parse out the remaining markdown content into chunks.

        Args:
            markdown (str): The markdown text containing the remaining content.
            tables (list[MarkdownTable]): The list of tables.
            code (list[MarkdownCode]): The list of code snippets.
            images (list[MarkdownImage]): The list of images.

        Returns:
            list[Chunk]: The list of chunks.

        """
        chunks: list[Chunk] = []

        # Get all the occupied
        occupied_indices: list[tuple[int, int]] = []
        occupied_indices.extend([(table.start_index, table.end_index) for table in tables])
        occupied_indices.extend([(code.start_index, code.end_index) for code in code])
        occupied_indices.extend([(image.start_index, image.end_index) for image in images])

        # Sort the occupied indices, by start and end index
        occupied_indices.sort(key=lambda x: (x[0], x[1]))

        # Get the remaining indices
        current_index = 0
        remaining_indices: list[tuple[int, int]] = []
        for index in occupied_indices:
            if index[0] > current_index:
                remaining_indices.append((current_index, index[0]))
            current_index = index[1]
        if current_index < len(markdown):
            remaining_indices.append((current_index, len(markdown)))

        # Get the chunks
        for index in remaining_indices:
            # Start and end index
            start_index = index[0]
            end_index = index[1]
            text = markdown[start_index:end_index]

            # Only create chunk if it contains meaningful content (not just whitespace)
            if text.strip():
                token_count = self.tokenizer.count_tokens(text)
                chunks.append(
                    Chunk(
                        text=text,
                        start_index=start_index,
                        end_index=end_index,
                        token_count=token_count,
                    ),
                )

        return chunks

    def parse(self, text: str) -> MarkdownDocument:
        """Parse markdown text directly into a MarkdownDocument.

        Args:
            text (str): The markdown text to parse.

        Returns:
            MarkdownDocument: The processed markdown document.

        """
        logger.debug(f"Processing markdown text: {len(text)} characters")

        # Extract all the tables, code snippets, and images
        tables = self.prepare_tables(text)
        code = self.prepare_code(text)
        images = self.extract_images(text)

        # Extract the chunks
        chunks: list[Chunk] = self.extract_chunks(text, tables, code, images)

        logger.info(
            f"Markdown processing complete: extracted {len(tables)} tables, {len(code)} code blocks, {len(images)} images, {len(chunks)} chunks",
        )
        return MarkdownDocument(
            content=text,
            tables=tables,
            code=code,
            images=images,
            chunks=chunks,
        )

    def process(self, path: str | os.PathLike) -> MarkdownDocument:
        """Process a markdown file into a MarkdownDocument.

        Args:
            path: The path to the markdown file.

        Returns:
            MarkdownDocument: The processed markdown document.

        """
        # Read the markdown file
        markdown = self.read(path)

        # Use parse to process the content
        return self.parse(markdown)
