from __future__ import annotations

import os
from collections.abc import Iterable
from pathlib import Path

import yaml


class PromptLoader:
    def __init__(self, prompts_path: Path | Iterable[Path] | None = None) -> None:
        self._paths = self._resolve_paths(prompts_path)
        self._prompts = self._load_prompts()

    def load(self, name: str) -> str:
        return str(self._prompts.get(name, "") or "").strip()

    def _load_prompts(self) -> dict[str, str]:
        merged: dict[str, str] = {}
        for path in self._paths:
            if not path.is_file():
                continue
            payload = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
            if not isinstance(payload, dict):
                continue
            namespace = str(payload.get("namespace") or "").strip()
            prompts = payload.get("prompts", payload)
            if not isinstance(prompts, dict):
                continue
            for key, value in prompts.items():
                prompt_name = f"{namespace}.{key}" if namespace else str(key)
                merged[prompt_name] = str(value or "")
        return merged

    def _resolve_paths(self, prompts_path: Path | Iterable[Path] | None) -> tuple[Path, ...]:
        if prompts_path is None:
            base = Path(__file__).resolve().parent / "prompts.yml"
            env_override = os.getenv("AGENT_PROMPTS_DIR", "").strip()
            raw_path = Path(env_override) if env_override else base
            return (raw_path / "prompts.yml" if raw_path.is_dir() else raw_path,)
        if isinstance(prompts_path, Path):
            return (prompts_path,)
        return tuple(Path(item) for item in prompts_path)
