"""Suite loaders and schema helpers for the quality framework."""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import yaml

from logos_adk.errors import ConfigError
from logos_adk.testing_framework import (
    GoldenExample,
    IntegrationScenario,
    IntegrationStep,
    JudgeCase,
    LoadTestConfig,
)


SCHEMA_DIR = Path(__file__).with_name("schemas") / "quality"
SCHEMA_FILES = {
    "integration": SCHEMA_DIR / "integration_suite.schema.json",
    "load-test": SCHEMA_DIR / "load_test_suite.schema.json",
    "regression": SCHEMA_DIR / "regression_suite.schema.json",
    "judge-eval": SCHEMA_DIR / "judge_eval_suite.schema.json",
}

SUITE_TEMPLATES: dict[str, dict[str, Any]] = {
    "integration": {
        "scenarios": [
            {
                "name": "multi_turn_context",
                "steps": [
                    {"prompt": "My name is Alex", "session_id": "demo-session"},
                    {
                        "prompt": "What is my name?",
                        "session_id": "demo-session",
                        "expect_contains": ["Alex"],
                    },
                ],
            }
        ]
    },
    "load-test": {
        "prompts": ["hello", "summarize this", "give me a plan"],
        "concurrency": 4,
        "iterations": 5,
        "session_prefix": "load",
    },
    "regression": {
        "examples": [
            {"name": "exact_match", "prompt": "Say hello", "expected": "Hello", "mode": "exact"},
            {
                "name": "contains_match",
                "prompt": "Explain observability",
                "expected": "metrics",
                "mode": "contains",
            },
        ]
    },
    "judge-eval": {
        "judge": {
            "pass_threshold": 0.7,
            "mock_response": '{"score": 0.9, "rationale": "good answer"}',
        },
        "cases": [
            {
                "name": "quality_case",
                "prompt": "Explain why PostgreSQL fits production workloads",
                "criteria": "Answer should be technically accurate and mention concurrency or durability",
                "reference": "PostgreSQL supports concurrent production workloads and durable storage",
            }
        ],
    },
}


def load_suite_file(path: str | Path) -> dict[str, Any]:
    """Load a YAML or JSON suite file."""
    suite_path = Path(path)
    content = suite_path.read_text()
    if suite_path.suffix.lower() == ".json":
        data = json.loads(content)
    else:
        data = yaml.safe_load(content)
    if not isinstance(data, dict):
        raise ConfigError("Suite file must contain a top-level mapping")
    return data


def load_schema(kind: str) -> dict[str, Any]:
    """Load the JSON schema for a suite kind."""
    try:
        path = SCHEMA_FILES[kind]
    except KeyError as exc:
        valid = ", ".join(sorted(SCHEMA_FILES))
        raise ConfigError(f"Unknown quality schema kind: {kind}. Valid values: {valid}") from exc
    return json.loads(path.read_text())


def render_suite_template(kind: str, *, format: str = "yaml") -> str:
    """Render a starter suite template in YAML or JSON."""
    try:
        payload = SUITE_TEMPLATES[kind]
    except KeyError as exc:
        valid = ", ".join(sorted(SUITE_TEMPLATES))
        raise ConfigError(f"Unknown quality template kind: {kind}. Valid values: {valid}") from exc
    if format == "json":
        return json.dumps(payload, indent=2, ensure_ascii=False)
    if format != "yaml":
        raise ConfigError("Template format must be 'yaml' or 'json'")
    return yaml.safe_dump(payload, sort_keys=False, allow_unicode=True)


def _expect_list(data: dict[str, Any], key: str) -> list[Any]:
    value = data.get(key)
    if not isinstance(value, list):
        raise ConfigError(f"Suite requires '{key}' to be a list")
    return value


def parse_integration_suite(data: dict[str, Any]) -> list[IntegrationScenario]:
    """Parse integration scenarios from suite data."""
    scenarios: list[IntegrationScenario] = []
    for scenario in _expect_list(data, "scenarios"):
        if not isinstance(scenario, dict):
            raise ConfigError("integration.scenarios items must be mappings")
        steps_data = scenario.get("steps")
        if not isinstance(steps_data, list) or not steps_data:
            raise ConfigError("integration scenario requires a non-empty 'steps' list")
        steps = []
        for step in steps_data:
            if not isinstance(step, dict):
                raise ConfigError("integration.steps items must be mappings")
            prompt = step.get("prompt")
            if not isinstance(prompt, str) or not prompt:
                raise ConfigError("integration step requires non-empty 'prompt'")
            steps.append(
                IntegrationStep(
                    prompt=prompt,
                    session_id=step.get("session_id"),
                    expect_contains=list(step.get("expect_contains", []) or []),
                    expect_not_contains=list(step.get("expect_not_contains", []) or []),
                )
            )
        scenarios.append(
            IntegrationScenario(
                name=str(scenario.get("name", f"scenario_{len(scenarios) + 1}")),
                steps=steps,
            )
        )
    return scenarios


def parse_load_test_suite(data: dict[str, Any]) -> LoadTestConfig:
    """Parse load test config from suite data."""
    prompts = _expect_list(data, "prompts")
    normalized_prompts = [prompt for prompt in prompts if isinstance(prompt, str) and prompt.strip()]
    if not normalized_prompts:
        raise ConfigError("load-test suite requires at least one string prompt")
    return LoadTestConfig(
        prompts=normalized_prompts,
        concurrency=int(data.get("concurrency", 4)),
        iterations=int(data.get("iterations", 1)),
        session_prefix=str(data.get("session_prefix", "load")),
    )


def parse_regression_suite(data: dict[str, Any]) -> list[GoldenExample]:
    """Parse regression/golden dataset suite."""
    examples: list[GoldenExample] = []
    for item in _expect_list(data, "examples"):
        if not isinstance(item, dict):
            raise ConfigError("regression.examples items must be mappings")
        prompt = item.get("prompt")
        expected = item.get("expected")
        if not isinstance(prompt, str) or not isinstance(expected, str):
            raise ConfigError("regression example requires string 'prompt' and 'expected'")
        examples.append(
            GoldenExample(
                name=str(item.get("name", f"example_{len(examples) + 1}")),
                prompt=prompt,
                expected=expected,
                session_id=item.get("session_id"),
                mode=str(item.get("mode", "exact")),
            )
        )
    return examples


def parse_judge_eval_suite(data: dict[str, Any]) -> tuple[list[JudgeCase], dict[str, Any]]:
    """Parse judge-eval cases and judge metadata."""
    cases: list[JudgeCase] = []
    for item in _expect_list(data, "cases"):
        if not isinstance(item, dict):
            raise ConfigError("judge-eval.cases items must be mappings")
        prompt = item.get("prompt")
        criteria = item.get("criteria")
        if not isinstance(prompt, str) or not isinstance(criteria, str):
            raise ConfigError("judge-eval case requires string 'prompt' and 'criteria'")
        cases.append(
            JudgeCase(
                name=str(item.get("name", f"case_{len(cases) + 1}")),
                prompt=prompt,
                criteria=criteria,
                reference=item.get("reference"),
                session_id=item.get("session_id"),
            )
        )
    judge_config = data.get("judge", {})
    if judge_config is None:
        judge_config = {}
    if not isinstance(judge_config, dict):
        raise ConfigError("judge-eval suite 'judge' must be a mapping")
    return cases, judge_config


__all__ = [
    "SCHEMA_DIR",
    "SCHEMA_FILES",
    "load_schema",
    "load_suite_file",
    "parse_integration_suite",
    "parse_judge_eval_suite",
    "parse_load_test_suite",
    "parse_regression_suite",
    "render_suite_template",
    "SUITE_TEMPLATES",
]
