from __future__ import annotations

from pathlib import Path

import yaml

from tests.pipeline_setup_v4.core.models import CaseExpectations, RetrievalPlanExpectation, RouterExpectation, V4Case


class CaseDirectoryLoader:
    def load(self, cases_dir: Path) -> list[V4Case]:
        if cases_dir.is_file():
            return self._load_file(cases_dir)
        files = sorted(path for path in cases_dir.rglob("*.yaml") if path.is_file())
        if not files:
            raise ValueError(f"No YAML case files found in: {cases_dir}")
        cases: list[V4Case] = []
        for path in files:
            cases.extend(self._load_file(path))
        return cases

    def _load_file(self, path: Path) -> list[V4Case]:
        payload = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
        if not isinstance(payload, dict):
            raise ValueError(f"Invalid case file {path}: expected mapping")
        defaults = dict(payload.get("defaults") or {})
        items = payload.get("cases") or []
        if not isinstance(items, list):
            raise ValueError(f"Invalid case file {path}: `cases` must be a list")
        return [self._to_case(path, raw, defaults) for raw in items]

    def _to_case(self, path: Path, raw: object, defaults: dict) -> V4Case:
        if not isinstance(raw, dict):
            raise ValueError(f"Invalid case in {path}: expected object")
        case_id = str(raw.get("id") or "").strip()
        component = str(raw.get("component") or defaults.get("component") or "").strip()
        query = str(raw.get("query") or "").strip()
        rag_session_id = str(raw.get("rag_session_id") or defaults.get("rag_session_id") or "").strip() or None
        route = dict(raw.get("route") or {})
        if not route and isinstance(defaults.get("route"), dict):
            route = dict(defaults.get("route") or {})
        if not case_id or not component:
            raise ValueError(f"Invalid case in {path}: `id` and `component` are required")
        if component in {
            "process_v2_intent_router",
            "process_v2_router_plus_retrieval_policy",
            "process_v2_router_plus_retrieval_policy_rag",
            "process_v2_full_chain",
        } and not query:
            raise ValueError(f"Invalid case in {path}: `query` is required for {component}")
        if component == "process_v2_retrieval_policy_resolver" and not route:
            raise ValueError(f"Invalid case in {path}: `route` is required for {component}")
        expected = dict(raw.get("expected") or {})
        return V4Case(
            case_id=case_id,
            component=component,  # type: ignore[arg-type]
            query=query,
            rag_session_id=rag_session_id,
            route=route,
            source_file=path,
            expectations=self._to_expectations(expected),
            notes=str(raw.get("notes") or ""),
            tags=tuple(str(item).strip() for item in raw.get("tags") or [] if str(item).strip()),
        )

    def _to_expectations(self, raw: dict) -> CaseExpectations:
        router = dict(raw.get("router") or {})
        route = dict(raw.get("route") or {})
        retrieval_plan = dict(raw.get("retrieval_plan") or raw.get("plan") or {})
        rag = dict(raw.get("rag") or {})
        pipeline = dict(raw.get("pipeline") or {})
        llm = dict(raw.get("llm") or {})
        return CaseExpectations(
            router=RouterExpectation(
                domain=str(router.get("domain") or "").strip() or None,
                intent=str(router.get("intent") or "").strip() or None,
                sub_intent=str(router.get("sub_intent") or "").strip() or None,
            ),
            retrieval_plan=RetrievalPlanExpectation(
                profile=str(retrieval_plan.get("profile") or "").strip() or None,
                layers=tuple(str(item).strip() for item in retrieval_plan.get("layers") or [] if str(item).strip()),
                limit=int(retrieval_plan["limit"]) if retrieval_plan.get("limit") is not None else None,
                filters=self._plain_mapping(dict(retrieval_plan.get("filters") or {})),
            ),
            route_assertions=route,
            retrieval_plan_assertions=retrieval_plan,
            rag_assertions=rag,
            pipeline_assertions=pipeline,
            llm_assertions=llm,
        )

    def _plain_mapping(self, raw: dict[str, object]) -> dict[str, object]:
        plain: dict[str, object] = {}
        for key, value in raw.items():
            if self._is_assertion_key(key) or value in {"present", "absent"}:
                continue
            plain[key] = value
        return plain

    def _is_assertion_key(self, key: str) -> bool:
        suffixes = ("_not_contains", "_contains_any", "_contains", "_equals_any", "_one_of")
        return any(key.endswith(suffix) for suffix in suffixes)
