from __future__ import annotations

from pathlib import Path
from typing import List, Tuple

from pydantic import BaseModel, field_validator
from pydantic_yaml import parse_yaml_file_as
from ruamel.yaml import YAML

from srs2std.srs.structure import (
    SectionTitle,
    TableOfContents,
    build_toc_from_md,
    split_md_by_toc,
)
from srs2std.util import clean_md_to_text, find_and_split_by_regex, sent_tokenize


class SRS(BaseModel):
    sections: List[Section]

    def to_str(self, ignore_title: bool = True, ignore_headers: bool = True) -> str:
        return "\n\n".join(
            [section.to_str(ignore_title, ignore_headers) for section in self.sections]
        )

    def to_sentences(
        self, ignore_title: bool = True, ignore_headers: bool = True
    ) -> List[str]:
        sentences = []
        for section in self.sections:
            sentences.extend(section.to_sentences(ignore_title, ignore_headers))
        return sentences


class Section(BaseModel):
    title: SectionTitle
    content: str
    subsections: List[Subsection]

    @field_validator("content", mode="before")
    def strip_content(cls, v: str) -> str:
        return v.strip()

    def to_str(self, ignore_title: bool = True, ignore_headers: bool = True) -> str:
        texts = []
        for line in self.content.splitlines():
            if "[N/A" in line or "[NULL" in line:
                continue
            texts.append(line)

        for subsection in self.subsections:
            subsection_content = subsection.to_str(ignore_headers)
            if subsection_content:
                texts.append(subsection_content)

        text = "\n".join(texts).strip()
        if not ignore_title and text:
            text = "\n".join([self.title.to_str(), text])
        return text

    def to_sentences(
        self, ignore_title: bool = True, ignore_headers: bool = True
    ) -> List[str]:
        cleaned_text = clean_md_to_text(self.to_str(ignore_title, ignore_headers))
        return sent_tokenize(cleaned_text)


class Subsection(BaseModel):
    header: str
    content: str
    subsections: List[Subsection]

    @field_validator("header", mode="before")
    def strip_header(cls, v: str) -> str:
        return v.strip()

    @field_validator("content", mode="before")
    def strip_content(cls, v: str) -> str:
        return v.strip()

    def to_str(self, ignore_headers: bool = True) -> str:
        texts = []
        for line in self.content.splitlines():
            if "[N/A" in line or "[NULL" in line:
                continue
            texts.append(line)

        for subsection in self.subsections:
            subsection_content = subsection.to_str(ignore_headers)
            if subsection_content:
                texts.append(subsection_content)

        text = "\n".join(texts).strip()
        if not ignore_headers and text:
            text = "\n".join([self.header, text])
        return text


def parse_md_as_srs(md_doc: str, toc: TableOfContents | None = None) -> SRS:
    toc = toc or build_toc_from_md(md_doc)
    section_titles, section_contents = split_md_by_toc(md_doc, toc)
    sections = []
    for title, content in zip(section_titles, section_contents):
        section = parse_content_as_section(title, content)
        sections.append(section)
    return SRS(sections=sections)


def parse_content_as_section(title: SectionTitle, content: str) -> Section:
    headers_pattern = r"#+\s+.*"
    headers, contents = find_and_split_by_regex(headers_pattern, content)
    if headers:
        subsections = parse_subsections(headers, contents[1:])
        return Section(title=title, content=contents[0], subsections=subsections)
    else:
        return Section(title=title, content=contents[0], subsections=[])


def parse_subsections(headers: List[str], contents: List[str]) -> List[Subsection]:
    def _parse_subsections(
        headers: List[str],
        contents: List[str],
    ) -> Tuple[List[Subsection], List[str], List[str]]:
        if not headers:
            return [], [], []

        base_level = headers[0].count("#")
        subsections = []

        while headers:
            current_level = headers[0].count("#")
            if current_level != base_level:
                break

            subsection = Subsection(
                header=headers[0],
                content=contents[0],
                subsections=[],
            )
            headers, contents = headers[1:], contents[1:]

            if headers and headers[0].count("#") > base_level:
                nested, headers, contents = _parse_subsections(headers, contents)
                subsection.subsections = nested
            subsections.append(subsection)

        return subsections, headers, contents

    subsections, _, _ = _parse_subsections(headers, contents)
    return subsections


def load_srs_md(path: Path, toc: TableOfContents | None = None) -> SRS:
    with path.open("r", encoding="utf-8") as f:
        md_doc = f.read()
    return parse_md_as_srs(md_doc, toc)


def load_srs_yaml(path: Path) -> SRS:
    return parse_yaml_file_as(SRS, path)


def save_srs_yaml(srs: SRS, path: Path):
    def _str_representer(dumper, data):
        if isinstance(data, str) and "\n" in data:
            return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
        return dumper.represent_scalar("tag:yaml.org,2002:str", data)

    yaml = YAML()
    yaml.default_flow_style = False
    yaml.indent(sequence=4, offset=2)
    yaml.representer.add_representer(str, _str_representer)

    with path.open("w", encoding="utf-8") as f:
        yaml.dump(srs.model_dump(), f)
