from __future__ import annotations

import uuid
from pathlib import Path
from typing import Any

import yaml
from pydantic import BaseModel, ValidationError
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from core.db.session import AsyncSessionLocal
from core.logging import get_logger
from models.llm import Llm
from models.llm_factory import LlmFactory

logger = get_logger("core.config.initial.init_data")


class LlmFactorySeed(BaseModel):
    name: str
    request_url: str
    avatar: str | None = None


class LlmSeed(BaseModel):
    model_code: str
    factory_name: str


class LlmCatalogSeed(BaseModel):
    factories: list[LlmFactorySeed]
    llms: list[LlmSeed]


def _default_catalog_path() -> Path:
    return (
        Path(__file__).resolve().parents[1] / "static" / "database" / "llm_catalog.yaml"
    )


def load_llm_catalog(catalog_path: Path | None = None) -> dict[str, Any]:
    path = catalog_path or _default_catalog_path()
    with path.open("r", encoding="utf-8") as file:
        loaded = yaml.safe_load(file) or {}
    if not isinstance(loaded, dict):
        raise ValueError(f"Invalid LLM catalog format: {path}")
    raw_factories = loaded.get("factories", [])
    raw_llms = loaded.get("llms", [])
    if not isinstance(raw_factories, list) or not isinstance(raw_llms, list):
        raise ValueError(f"Invalid LLM catalog sections: {path}")
    try:
        parsed = LlmCatalogSeed.model_validate(
            {
                "factories": list(raw_factories),
                "llms": list(raw_llms),
            }
        )
    except ValidationError as exc:
        raise ValueError(f"Invalid LLM catalog data: {path}") from exc

    return parsed.model_dump()


async def _upsert_factory(
    session: AsyncSession,
    *,
    name: str,
    request_url: str,
    avatar: str | None,
) -> uuid.UUID:
    result = await session.execute(select(LlmFactory).where(LlmFactory.name == name))
    factory = result.scalar_one_or_none()

    if factory is None:
        factory = LlmFactory(name=name, request_url=request_url, avatar=avatar)
        session.add(factory)
        await session.flush()
    else:
        factory.request_url = request_url
        factory.avatar = avatar

    return factory.id


async def _upsert_llm(
    session: AsyncSession,
    *,
    model_code: str,
    factory_id: uuid.UUID,
) -> None:
    result = await session.execute(select(Llm).where(Llm.model_code == model_code))
    llm = result.scalar_one_or_none()
    if llm is None:
        session.add(Llm(model_code=model_code, factory_id=factory_id))
        return
    llm.factory_id = factory_id


async def initialize_data() -> bool:
    """Initialize bootstrap data."""
    catalog = load_llm_catalog()

    async with AsyncSessionLocal() as session:
        async with session.begin():
            factory_id_by_name: dict[str, uuid.UUID] = {}
            for factory in catalog["factories"]:
                factory_id = await _upsert_factory(
                    session,
                    name=factory["name"],
                    request_url=factory["request_url"],
                    avatar=factory.get("avatar"),
                )
                factory_id_by_name[factory["name"]] = factory_id

            for llm in catalog["llms"]:
                factory_name = llm["factory_name"]
                resolved_factory_id = factory_id_by_name.get(factory_name)
                if resolved_factory_id is None:
                    raise RuntimeError(
                        f"Factory '{factory_name}' not found for model '{llm['model_code']}'"
                    )
                await _upsert_llm(
                    session,
                    model_code=llm["model_code"],
                    factory_id=resolved_factory_id,
                )

    logger.info("Initialized LLM factory/model seed data")
    return True
