import re
from pathlib import Path
from typing import Literal

import yaml
from pydantic import BaseModel, ConfigDict, Field, model_validator


class StrictModel(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)


class SplitConfig(StrictModel):
    families: int = Field(gt=0)
    depths: list[int]
    motifs: list[Literal["chain", "conjunction", "relational", "diamond"]]
    heldout_wording: bool

    @model_validator(mode="after")
    def validate_design(self):
        if not self.depths or min(self.depths) < 0 or max(self.depths) > 5:
            raise ValueError("Reasoning depths must be nonempty and between zero and five")
        if not self.motifs:
            raise ValueError("At least one motif is required")
        if "diamond" in self.motifs and self.depths != [3]:
            raise ValueError("The reserved diamond motif requires depth three")
        return self


class DatasetConfig(StrictModel):
    schema_version: Literal[1]
    seed: int
    max_attempts_per_family: int = Field(gt=0)
    entities_min: int = Field(ge=3)
    entities_max: int
    extra_facts_min: int = Field(ge=0)
    extra_facts_max: int
    max_text_characters: int = Field(gt=0)
    train_templates: list[Literal["canonical", "conditional"]]
    heldout_template: Literal["paraphrase"]
    train_entities: list[str]
    heldout_entities: list[str]
    predicates: list[list[str]]
    relations: list[str]
    splits: dict[str, SplitConfig]

    @model_validator(mode="after")
    def validate_design(self):
        if self.entities_max < self.entities_min:
            raise ValueError("entities_max must be at least entities_min")
        if self.extra_facts_max < self.extra_facts_min:
            raise ValueError("extra_facts_max must be at least extra_facts_min")
        if not self.train_templates or not self.relations or not self.splits:
            raise ValueError("Templates, relations, and splits must be nonempty")
        if set(self.train_entities) & set(self.heldout_entities):
            raise ValueError("Held-out entity names overlap with training entity names")
        for names in (self.train_entities, self.heldout_entities):
            if len(set(names)) != len(names) or len(names) < self.entities_max:
                raise ValueError("Entity pools must be unique and cover entities_max")
        if not self.predicates or any(len(set(pool)) < 8 for pool in self.predicates):
            raise ValueError("Each predicate vocabulary needs eight unique predicates")
        words = [*self.relations, *(word for pool in self.predicates for word in pool)]
        if any(re.fullmatch(r"[a-z][a-z_]*", word) is None for word in words):
            raise ValueError("Predicate names must be lowercase identifiers")
        if set(self.relations) & {word for pool in self.predicates for word in pool}:
            raise ValueError("Unary and binary predicates cannot share a name")
        if "train" not in self.splits:
            raise ValueError("A training split is required")
        train = self.splits["train"]
        if train.heldout_wording or "diamond" in train.motifs or max(train.depths) > 3:
            raise ValueError("Training cannot contain held-out wording, motifs, or depths")
        return self


def load_dataset_config(path: Path) -> DatasetConfig:
    with path.open() as stream:
        return DatasetConfig.model_validate(yaml.safe_load(stream))