"""Loads the schema for the vaultgen agent."""

import json
import re
from datetime import date, datetime
from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, Field, create_model


def load_raw_schema(vault: Path) -> dict:
    """Load schema from the ```json code block in config.md (JSON)."""
    path = vault.joinpath("config.md")
    text = path.read_text()
    match = re.search(r"```json\s*\n(.*?)```", text, re.DOTALL)
    if not match:
        raise ValueError(f"No ```json``` block found in {path}")
    return json.loads(match.group(1).strip())


def get_type(type_str: str) -> Any:
    """Get the type from the string."""
    if isinstance(type_str, list):
        return Literal[*type_str]

    type_map: dict[str, type] = {
        "datetime": datetime,
        "date": date,
        "link": str,
        "list": list,
        "str": str,
        "int": int,
        "float": float,
        "bool": bool,
    }

    if type_str in type_map:
        return type_map[type_str]

    raise ValueError(f"Unknown type: {type_str}")


def load_schema(vault: Path) -> type[BaseModel]:
    """Load schema from the ```json code block in config.md (JSON)."""
    schema = load_raw_schema(vault)
    fields = {}
    for key, value in schema.items():
        field_type = get_type(value["type"])
        if value.get("optional", False):
            field_type = field_type | None

        description = value["description"]
        if value["type"] == "link":
            description += " (it should be the exact name of the linked note in double brackets: e.g. '[[LINKED NOTE NAME]]')"

        fields[key] = (
            field_type,
            Field(description=description),
        )

    fields["name"] = (
        str,
        Field(description="The name of the file and title of the note"),
    )
    fields["content"] = (
        str | None,
        Field(description="The content of the note", default=None),
    )

    docstring = "The schema for the vaultgen agent."
    return create_model("Schema", **fields, __doc__=docstring)
