# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal

import regex as re
import yaml

logger = logging.getLogger("mkdocs")

ROOT_DIR = Path(__file__).parent.parent.parent.parent
ROOT_DIR_RELATIVE = "../../../../.."
EXAMPLE_DIR = ROOT_DIR / "examples"
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
NAV_FILE = ROOT_DIR / "docs/.nav.yml"


def fix_case(text: str) -> str:
    subs = {
        "api": "API",
        "cli": "CLI",
        "cpu": "CPU",
        "llm": "LLM",
        "mae": "MAE",
        "tpu": "TPU",
        "gguf": "GGUF",
        "lora": "LoRA",
        "rlhf": "RLHF",
        "vllm": "vLLM",
        "openai": "OpenAI",
        "lmcache": "LMCache",
        "multilora": "MultiLoRA",
        "mlpspeculator": "MLPSpeculator",
        r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
    }
    for pattern, repl in subs.items():
        text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
    return text


@dataclass
class Example:
    """
    Example class for generating documentation content from a given path.

    Attributes:
        path (Path): The path to the main directory or file.
        category (str): The category of the document.
        main_file (Path): The main file in the directory.
        other_files (list[Path]): list of other files in the directory.
        title (str): The title of the document.

    Methods:
        __post_init__(): Initializes the main_file, other_files, and title attributes.
        determine_main_file() -> Path: Determines the main file in the given path.
        determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
        determine_title() -> str: Determines the title of the document.
        generate() -> str: Generates the documentation content.
    """  # noqa: E501

    path: Path
    category: str = None
    main_file: Path = field(init=False)
    other_files: list[Path] = field(init=False)
    title: str = field(init=False)

    def __post_init__(self):
        self.main_file = self.determine_main_file()
        self.other_files = self.determine_other_files()
        self.title = self.determine_title()

    @property
    def is_code(self) -> bool:
        return self.main_file.suffix != ".md"

    def determine_main_file(self) -> Path:
        """
        Determines the main file in the given path.
        If the path is a file, it returns the path itself. Otherwise, it searches
        for Markdown files (*.md) in the directory and returns the first one found.
        Returns:
            Path: The main file path, either the original path if it's a file or the first
            Markdown file found in the directory.
        Raises:
            IndexError: If no Markdown files are found in the directory.
        """  # noqa: E501
        return self.path if self.path.is_file() else list(self.path.glob("*.md")).pop()

    def determine_other_files(self) -> list[Path]:
        """
        Determine other files in the directory excluding the main file.

        This method checks if the given path is a file. If it is, it returns an empty list.
        Otherwise, it recursively searches through the directory and returns a list of all
        files that are not the main file.

        Returns:
            list[Path]: A list of Path objects representing the other files in the directory.
        """  # noqa: E501
        if self.path.is_file():
            return []
        # Binary file extensions to exclude
        binary_extensions = {
            ".wav",
            ".mp3",
            ".mp4",
            ".avi",
            ".mov",
            ".mkv",  # Audio/Video
            ".png",
            ".jpg",
            ".jpeg",
            ".gif",
            ".bmp",
            ".ico",
            ".svg",  # Images
            ".pdf",
            ".zip",
            ".tar",
            ".gz",
            ".bz2",
            ".xz",  # Archives/Documents
            ".exe",
            ".so",
            ".dll",
            ".dylib",  # Binaries
            ".bin",
            ".dat",
            ".db",
            ".sqlite",  # Data files
        }

        def is_other_file(file: Path) -> bool:
            return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions

        return [file for file in self.path.rglob("*") if is_other_file(file)]

    def determine_title(self) -> str:
        if not self.is_code:
            # Specify encoding for building on Windows
            with open(self.main_file, encoding="utf-8") as f:
                first_line = f.readline().strip()
            match = re.match(r"^#\s+(?P<title>.+)$", first_line)
            if match:
                return match.group("title")
        return fix_case(self.path.stem.replace("_", " ").title())

    def fix_relative_links(self, content: str) -> str:
        """
        Fix relative links in markdown content by converting them to gh-file
        format.

        Args:
            content (str): The markdown content to process

        Returns:
            str: Content with relative links converted to gh-file format
        """
        # Regex to match markdown links [text](relative_path)
        # This matches links that don't start with http, https, ftp, or #
        link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"

        def replace_link(match):
            link_text = match.group(1)
            relative_path = match.group(2)

            # Make relative to repo root
            gh_file = (self.main_file.parent / relative_path).resolve()
            gh_file = gh_file.relative_to(ROOT_DIR)

            # Make GitHub URL
            url = "https://github.com/vllm-project/vllm-omni/"
            url += "tree/main" if self.path.is_dir() else "blob/main"
            gh_url = f"{url}/{gh_file}"

            return f"[{link_text}]({gh_url})"

        return re.sub(link_pattern, replace_link, content)

    def generate(self) -> str:
        content = f"# {self.title}\n\n"
        url = "https://github.com/vllm-project/vllm-omni/"
        url += "tree/main" if self.path.is_dir() else "blob/main"
        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"

        # Use long code fence to avoid issues with
        # included files containing code fences too
        code_fence = "``````"

        if self.is_code:
            main_file_rel = self.main_file.relative_to(ROOT_DIR)
            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
        else:
            with open(self.main_file, encoding="utf-8") as f:
                # Skip the title from md snippets as it's been included above
                main_content = f.readlines()[1:]
            content += self.fix_relative_links("".join(main_content))
        content += "\n"

        if not self.other_files:
            return content

        content += "## Example materials\n\n"
        for file in sorted(self.other_files):
            content += f'??? abstract "{file.relative_to(self.path)}"\n'
            if file.suffix != ".md":
                content += f"    {code_fence}{file.suffix[1:]}\n"
            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
            if file.suffix != ".md":
                content += f"    {code_fence}\n"

        return content


def update_nav_file(examples: list[Example]):
    """
    Update the .nav.yml file to include all generated examples.
    This function completely regenerates the examples section based on the actual
    folder structure, ensuring consistency between the examples folder and nav file.

    Args:
        examples: List of Example objects that have been generated
    """
    if not NAV_FILE.exists():
        logger.warning("Navigation file not found: %s", NAV_FILE)
        return

    # Read the current nav file
    with open(NAV_FILE, encoding="utf-8") as f:
        nav_data = yaml.safe_load(f) or {}

    nav_list = nav_data.get("nav", [])

    # Find the "User Guide" section
    user_guide_idx = None
    examples_idx = None
    for i, item in enumerate(nav_list):
        if isinstance(item, dict) and "User Guide" in item:
            user_guide_idx = i
            user_guide_content = item["User Guide"]
            # Find the "Examples" subsection
            for j, subitem in enumerate(user_guide_content):
                if isinstance(subitem, dict) and "Examples" in subitem:
                    examples_idx = j
                    break
            break

    if user_guide_idx is None or examples_idx is None:
        logger.warning("Could not find 'User Guide' -> 'Examples' section in nav file")
        return

    # Get existing Examples section to preserve non-example items (like README.md)
    existing_examples_content = nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"]

    # Preserve string items (like "examples/README.md") that are not example categories
    preserved_items = [
        item
        for item in existing_examples_content
        if isinstance(item, str) and not item.startswith("user_guide/examples/")
    ]

    # Group examples by category
    examples_by_category = {}
    for example in examples:
        category = example.category
        if category not in examples_by_category:
            examples_by_category[category] = []
        examples_by_category[category].append(example)

    # Build the new Examples section - start with preserved items
    examples_section = preserved_items.copy()

    # Add examples grouped by category, sorted by category name
    for category in sorted(examples_by_category.keys()):
        category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
        category_items = []
        for example in category_examples:
            doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
            rel_path = doc_path.relative_to(ROOT_DIR / "docs")
            category_items.append({example.title: str(rel_path)})

        if category_items:
            # Format category name (e.g., "offline_inference" -> "Offline Inference")
            category_title = fix_case(category.replace("_", " ").title())
            examples_section.append({category_title: category_items})

    # Update the nav structure
    nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section

    # Write back to file
    nav_data["nav"] = nav_list
    with open(NAV_FILE, "w", encoding="utf-8") as f:
        yaml.dump(nav_data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
    logger.info("Updated navigation file: %s", NAV_FILE.relative_to(ROOT_DIR))


def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
    logger.info("Generating example documentation")
    logger.debug("Root directory: %s", ROOT_DIR.resolve())
    logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
    logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve())

    # Create the EXAMPLE_DOC_DIR if it doesn't exist
    if not EXAMPLE_DOC_DIR.exists():
        EXAMPLE_DOC_DIR.mkdir(parents=True)

    categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir())

    examples = []
    glob_patterns = ["*.py", "*.md", "*.sh"]
    # Find categorised examples
    for category in categories:
        globs = [category.glob(pattern) for pattern in glob_patterns]
        for path in itertools.chain(*globs):
            examples.append(Example(path, category.stem))
        # Find examples in subdirectories
        for path in category.glob("*/*.md"):
            examples.append(Example(path.parent, category.stem))

    # Generate the example documentation
    for example in sorted(examples, key=lambda e: e.path.stem):
        example_name = f"{example.path.stem}.md"
        doc_path = EXAMPLE_DOC_DIR / example.category / example_name
        if not doc_path.parent.exists():
            doc_path.parent.mkdir(parents=True)
        # Specify encoding for building on Windows
        with open(doc_path, "w+", encoding="utf-8") as f:
            f.write(example.generate())
        logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR))

    # Update the navigation file
    update_nav_file(examples)