from __future__ import annotations

import re
from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

ListStyle = Literal["bullet", "numbered"]
ImageWidth = Literal["small", "medium", "large", "full"]
ImageAlignment = Literal["left", "center", "right"]
BlockType = Literal["heading", "paragraph", "list", "table", "callout", "quote", "image", "page_break"]

MAX_DOCUMENT_BLOCKS = 200
MAX_TABLE_COLUMNS = 8
MAX_TABLE_ROWS = 50
MAX_LIST_ITEMS = 50
_BLOCK_ID_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$")


class DocxSpecModel(BaseModel):
    """Strict base so unsupported LLM fields are repaired instead of silently discarded."""

    model_config = ConfigDict(extra="forbid")


def _non_empty(value: str, label: str) -> str:
    normalized = value.strip()
    if not normalized:
        raise ValueError(f"{label} must be non-empty")
    return normalized


class DocumentConfig(DocxSpecModel):
    title: str
    subtitle: str | None = None
    author: str | None = None
    date: str | None = None
    subject: str | None = None

    @model_validator(mode="after")
    def validate_text_fields(self) -> DocumentConfig:
        self.title = _non_empty(self.title, "Document title")
        for field_name in ("subtitle", "author", "date", "subject"):
            value = getattr(self, field_name)
            if value is not None:
                setattr(self, field_name, _non_empty(value, field_name.capitalize()))
        return self


class RichTextRun(DocxSpecModel):
    text: str
    bold: bool = False
    italic: bool = False
    underline: bool = False

    @model_validator(mode="after")
    def validate_text(self) -> RichTextRun:
        if not self.text:
            raise ValueError("Rich-text run text must be non-empty")
        return self


class BlockBase(DocxSpecModel):
    id: str

    @model_validator(mode="after")
    def validate_id(self) -> BlockBase:
        if not _BLOCK_ID_RE.fullmatch(self.id):
            raise ValueError(
                "Block id must start with a lowercase letter or digit and contain only lowercase letters, "
                "digits, underscores, or hyphens"
            )
        return self


class HeadingBlock(BlockBase):
    type: Literal["heading"]
    level: int = Field(..., ge=1, le=3)
    text: str

    @model_validator(mode="after")
    def validate_text(self) -> HeadingBlock:
        self.text = _non_empty(self.text, "Heading text")
        return self


class ParagraphBlock(BlockBase):
    type: Literal["paragraph"]
    text: str | None = None
    runs: list[RichTextRun] = Field(default_factory=list, max_length=100)

    @model_validator(mode="after")
    def validate_content(self) -> ParagraphBlock:
        if self.text is not None:
            self.text = _non_empty(self.text, "Paragraph text")
        if (self.text is None) == (not self.runs):
            raise ValueError("Paragraph requires exactly one of text or runs")
        if self.runs and not "".join(run.text for run in self.runs).strip():
            raise ValueError("Paragraph runs must contain visible text")
        return self


class ListItem(DocxSpecModel):
    text: str
    level: int = Field(default=0, ge=0, le=2)

    @model_validator(mode="after")
    def validate_text(self) -> ListItem:
        self.text = _non_empty(self.text, "List item text")
        return self


class ListBlock(BlockBase):
    type: Literal["list"]
    style: ListStyle = "bullet"
    items: list[ListItem] = Field(..., min_length=1, max_length=MAX_LIST_ITEMS)

    @model_validator(mode="after")
    def validate_level_progression(self) -> ListBlock:
        previous_level = 0
        for index, item in enumerate(self.items):
            if index == 0 and item.level != 0:
                raise ValueError("The first list item must be at level 0")
            if item.level > previous_level + 1:
                raise ValueError("List nesting cannot jump more than one level")
            previous_level = item.level
        return self


class TableBlock(BlockBase):
    type: Literal["table"]
    title: str | None = None
    columns: list[str] = Field(..., min_length=1, max_length=MAX_TABLE_COLUMNS)
    rows: list[list[str]] = Field(..., min_length=1, max_length=MAX_TABLE_ROWS)
    column_widths: list[float] | None = None
    source: str | None = None

    @model_validator(mode="after")
    def validate_table(self) -> TableBlock:
        if self.title is not None:
            self.title = _non_empty(self.title, "Table title")
        self.columns = [_non_empty(column, "Table column") for column in self.columns]
        if any(len(row) != len(self.columns) for row in self.rows):
            raise ValueError("Every table row must match the declared column count")
        if not any(any(str(cell).strip() for cell in row) for row in self.rows):
            raise ValueError("Table rows must contain at least one non-empty cell")
        self.rows = [[str(cell) for cell in row] for row in self.rows]
        if self.column_widths is not None:
            if len(self.column_widths) != len(self.columns):
                raise ValueError("column_widths must match the declared column count")
            if any(width <= 0 for width in self.column_widths):
                raise ValueError("column_widths values must be positive")
        if self.source is not None:
            self.source = _non_empty(self.source, "Table source")
        return self


class CalloutBlock(BlockBase):
    type: Literal["callout"]
    label: str | None = None
    text: str

    @model_validator(mode="after")
    def validate_content(self) -> CalloutBlock:
        if self.label is not None:
            self.label = _non_empty(self.label, "Callout label")
        self.text = _non_empty(self.text, "Callout text")
        return self


class QuoteBlock(BlockBase):
    type: Literal["quote"]
    text: str
    attribution: str | None = None

    @model_validator(mode="after")
    def validate_content(self) -> QuoteBlock:
        self.text = _non_empty(self.text, "Quote text")
        if self.attribution is not None:
            self.attribution = _non_empty(self.attribution, "Quote attribution")
        return self


class ImageBlock(BlockBase):
    type: Literal["image"]
    asset_id: str
    caption: str | None = None
    width: ImageWidth = "large"
    alignment: ImageAlignment = "center"

    @model_validator(mode="after")
    def validate_content(self) -> ImageBlock:
        self.asset_id = _non_empty(self.asset_id, "Image asset id")
        if self.caption is not None:
            self.caption = _non_empty(self.caption, "Image caption")
        return self


class PageBreakBlock(BlockBase):
    type: Literal["page_break"]


DocumentBlock = Annotated[
    HeadingBlock | ParagraphBlock | ListBlock | TableBlock | CalloutBlock | QuoteBlock | ImageBlock | PageBreakBlock,
    Field(discriminator="type"),
]


class DocumentSpec(DocxSpecModel):
    schema_version: Literal[1] = 1
    document: DocumentConfig
    blocks: list[DocumentBlock] = Field(..., min_length=1, max_length=MAX_DOCUMENT_BLOCKS)

    @model_validator(mode="after")
    def validate_document(self) -> DocumentSpec:
        block_ids = [block.id for block in self.blocks]
        if len(block_ids) != len(set(block_ids)):
            raise ValueError("Document block ids must be unique")
        if self.blocks[0].type == "page_break" or self.blocks[-1].type == "page_break":
            raise ValueError("Document cannot begin or end with a page break")
        for previous, current in zip(self.blocks, self.blocks[1:]):
            if previous.type == current.type == "page_break":
                raise ValueError("Document cannot contain consecutive page breaks")
        if not any(block.type not in {"heading", "page_break"} for block in self.blocks):
            raise ValueError("Document requires at least one substantive content block")
        return self
