from pathlib import Path
import logging
from typing import Any, Literal, TypeVar

import pydantic
import yaml  # type: ignore

T = TypeVar("T")

logger = logging.getLogger(__name__)
handler = logging.StreamHandler()
formatter = logging.Formatter(logging.BASIC_FORMAT)
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)


Url = str


class BaseModel(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(extra="forbid")


class OriginMetadata(BaseModel):
    name: str
    """Name of the ECS."""
    source: Url
    """Where the code is from."""
    upstream_source: None | Url
    """The original implementation."""
    paper: None | Url
    """The paper which introduces the system."""

    @pydantic.field_validator("name")
    @classmethod
    def check_name(cls, v: T, info: pydantic.ValidationInfo) -> T:
        if not isinstance(v, str):
            raise ValueError(f'Field "variants" must be a dict.')
        if info.context and (path := info.context.get("path", None)):
            if path.name != v:
                logger.warning(f"{path} does not match metadata name {v}")
        return v  # type: ignore


class SystemSpecificMetadata(BaseModel):
    game_type: Literal["signalling", "conversation", "navigation"]
    observation_type: Literal["vector", "image"]
    observation_continuous: bool
    data_source: None | Literal["natural", "synthetic"]
    variants: dict[str, Any]
    """Different variants of the environment.

    Each dict should be populated with the distinguishing characteristics of
    that environment.

    """
    game_subtype: None | str
    seeding_available: bool
    """System implements settable random seeds."""
    multi_step: bool
    """Environment has multiple timesteps per episode."""
    symmetric_agents: bool
    """Agents can send and receive messages."""
    multi_utterance: bool
    """Multiple utterances are present per line."""
    more_than_2_agents: bool
    """More than two agents present in the environment."""

    @pydantic.field_validator("variants")
    @classmethod
    def check_variants(cls, v: T, info: pydantic.ValidationInfo) -> T:
        if not isinstance(v, dict):
            raise ValueError(f'Field "variants" must be a dict.')
        if info.context and (path := info.context.get("path", None)):
            # Check that all listed variants are present on the filesystem.
            for k in v.keys():
                validate_variant(f"{path}/data/{k}")
            # Check that all variants on the filesystem are listed in the metadata.
            for var in Path(path).glob("systems/*"):
                if var.name not in v.keys():
                    logger.warning(
                        f'Directory "{path}/data/{var}" found but not listed in metadata.'
                    )
        return v  # type: ignore


def validate_variant(_path: str) -> None:
    path = Path(_path)
    if not path.exists():
        logger.warning(
            f"Variant {path.name} in metadata but not {path} does not exist."
        )
        return
    corpus_path = path / "corpus.jsonl"
    if not corpus_path.exists():
        logger.warning(f"Corpus file {corpus_path} does not exist.")
    metadata_path = path / "metadata.json"
    if not metadata_path.exists():
        logger.warning(f"Metadata file {metadata_path} does not exist.")


class SystemMetadata(BaseModel):
    origin: OriginMetadata
    system: SystemSpecificMetadata
    notes: str | None = None


def main() -> None:
    paths = list(Path(".").glob("systems/*"))
    for path in paths:
        try:
            # md_path = path / "metadata.yml"
            md_path = path / "system.json"
            if not md_path.exists():
                logger.warning(f'Missing metadata file at "{md_path}".')
                continue
            with md_path.open() as fo:
                raw = yaml.load(fo.read(), Loader=yaml.Loader)
            if raw is None:
                logger.warning(f'"{md_path}" is empty.')
                continue
            SystemMetadata.model_validate(raw, context={"path": path})
        except (yaml.parser.ParserError, pydantic.ValidationError) as e:
            logger.warning(f'Path "{md_path}" failed validation with:\n{e}')
    logger.info("Validation completed.")


if __name__ == "__main__":
    main()
