from __future__ import annotations

import re
from typing import List, Tuple

from pydantic import BaseModel, field_validator

from srs2std.util import find_and_split_by_regex


class TableOfContents(BaseModel):
    titles: List[SectionTitle]

    def overall_structure(self) -> str:
        return "\n".join([title.to_str() for title in self.titles])


class SectionTitle(BaseModel):
    number: str
    name: str

    @field_validator("number", mode="before")
    def validate_number(cls, v: str) -> str:
        v = v.strip()
        if not re.fullmatch(r"\d+(\.\d+)*", v):
            raise ValueError()
        return v

    @field_validator("name", mode="before")
    def capitalize_name(cls, v: str) -> str:
        return v.strip().capitalize()

    @classmethod
    def from_header(cls, header: str) -> SectionTitle:
        pattern = r"(\d+(?:\.\d+)*)\.?\s+(.+)"
        match = re.search(pattern, header, flags=re.I)
        number = match.group(1).strip() if match else ""
        name = match.group(2).strip() if match else header
        return SectionTitle(number=number, name=name)

    def to_str(self) -> str:
        level = self.number.count(".") + 1
        return " ".join(["#" * level, self.number, self.name])


def split_md_by_toc(
    md_doc: str, toc: TableOfContents
) -> Tuple[List[SectionTitle], List[str]]:
    titles_pattern = "|".join(
        rf"#+\s+{re.escape(title.number)}\.?\s+{re.escape(title.name)}\s*"
        for title in toc.titles
    )
    separators, separated = find_and_split_by_regex(titles_pattern, md_doc)
    section_titles = [SectionTitle.from_header(header) for header in separators]
    section_contents = separated[1:]
    return section_titles, section_contents


def build_toc_from_md(md_doc: str) -> TableOfContents:
    headers_pattern = r"#+\s+\d+(?:\.\d+)*\.?\s+.*"
    headers, _ = find_and_split_by_regex(headers_pattern, md_doc)
    titles = [SectionTitle.from_header(header) for header in headers]
    return TableOfContents(titles=titles)


def complement_falling_sections(md_doc: str, toc: TableOfContents) -> str:
    titles, contents = split_md_by_toc(md_doc, toc)

    # create new arrayes including missing section titles
    new_titles = []
    new_contents = []
    for toc_title in toc.titles:
        for title, content in zip(titles, contents):
            if toc_title.number == title.number and toc_title.name == title.name:
                new_contents.append(content or "[NULL]")
                break
        else:
            new_contents.append("[NULL]")
        new_titles.append(toc_title)

    # convert to markdown format
    fragments = []
    for title, content in zip(new_titles, new_contents):
        fragments.append(title.to_str())
        fragments.append(content)
        fragments.append("")
    return "\n".join(fragments)
