"""Integrated quality framework for Logos ADK."""
from __future__ import annotations

import asyncio
from contextlib import AbstractContextManager
from dataclasses import dataclass, field
from datetime import UTC, datetime
import json
import os
from pathlib import Path
import shutil
import statistics
import tempfile
from typing import Any, Callable

from logos_adk.agent import Agent
from logos_adk.observe import reset_observability_store
from logos_adk.runtime import _default_runtime_holder
from logos_adk.types import Response


@dataclass
class IntegrationStep:
    """Single integration step for an agent conversation."""

    prompt: str
    session_id: str | None = None
    expect_contains: list[str] = field(default_factory=list)
    expect_not_contains: list[str] = field(default_factory=list)


@dataclass
class IntegrationScenario:
    """Multi-step integration scenario."""

    name: str
    steps: list[IntegrationStep]


@dataclass
class IntegrationResult:
    """Result for a single integration scenario."""

    name: str
    passed: bool
    responses: list[Response] = field(default_factory=list)
    error: str | None = None


async def run_integration_suite(
    agent: Agent,
    scenarios: list[IntegrationScenario],
) -> list[IntegrationResult]:
    """Run multi-step integration scenarios against an agent."""
    results: list[IntegrationResult] = []
    for scenario in scenarios:
        responses: list[Response] = []
        try:
            for step in scenario.steps:
                response = await agent.run(step.prompt, session_id=step.session_id)
                responses.append(response)
                content = response.content
                for needle in step.expect_contains:
                    if needle not in content:
                        raise AssertionError(
                            f"Scenario {scenario.name!r} missing expected content {needle!r}"
                        )
                for needle in step.expect_not_contains:
                    if needle in content:
                        raise AssertionError(
                            f"Scenario {scenario.name!r} unexpectedly contained {needle!r}"
                        )
        except Exception as exc:
            results.append(
                IntegrationResult(
                    name=scenario.name,
                    passed=False,
                    responses=responses,
                    error=str(exc),
                )
            )
            continue
        results.append(IntegrationResult(name=scenario.name, passed=True, responses=responses))
    return results


@dataclass
class LoadTestConfig:
    """Configuration for a synthetic load test."""

    prompts: list[str]
    concurrency: int = 4
    iterations: int = 1
    session_prefix: str = "load"


@dataclass
class LoadTestReport:
    """Aggregated load test metrics."""

    total_requests: int
    success_count: int
    error_count: int
    avg_latency_ms: float
    p95_latency_ms: float
    max_latency_ms: float
    started_at: str
    finished_at: str
    errors: list[str] = field(default_factory=list)


async def run_load_test(agent: Agent, config: LoadTestConfig) -> LoadTestReport:
    """Run a concurrent synthetic load test against an agent."""
    semaphore = asyncio.Semaphore(max(1, config.concurrency))
    latencies: list[float] = []
    errors: list[str] = []
    started_at = datetime.now(UTC)

    requests: list[tuple[str, str | None]] = []
    for iteration in range(config.iterations):
        for index, prompt in enumerate(config.prompts):
            requests.append((prompt, f"{config.session_prefix}-{iteration}-{index}"))

    async def execute(prompt: str, session_id: str | None) -> None:
        started = asyncio.get_running_loop().time()
        try:
            async with semaphore:
                await agent.run(prompt, session_id=session_id)
        except Exception as exc:
            errors.append(str(exc))
        finally:
            finished = asyncio.get_running_loop().time()
            latencies.append((finished - started) * 1000)

    await asyncio.gather(*(execute(prompt, session_id) for prompt, session_id in requests))
    finished_at = datetime.now(UTC)
    if latencies:
        ordered = sorted(latencies)
        p95_index = max(0, min(len(ordered) - 1, round(len(ordered) * 0.95) - 1))
        p95 = ordered[p95_index]
        avg = statistics.fmean(ordered)
        max_latency = ordered[-1]
    else:
        p95 = 0.0
        avg = 0.0
        max_latency = 0.0
    return LoadTestReport(
        total_requests=len(requests),
        success_count=len(requests) - len(errors),
        error_count=len(errors),
        avg_latency_ms=round(avg, 2),
        p95_latency_ms=round(p95, 2),
        max_latency_ms=round(max_latency, 2),
        started_at=started_at.isoformat(),
        finished_at=finished_at.isoformat(),
        errors=errors,
    )


@dataclass
class GoldenExample:
    """Expected regression example."""

    name: str
    prompt: str
    expected: str
    session_id: str | None = None
    mode: str = "exact"


@dataclass
class GoldenResult:
    """Regression test result for a single golden example."""

    name: str
    passed: bool
    actual: str
    expected: str
    mode: str


async def run_golden_dataset(
    agent: Agent,
    examples: list[GoldenExample],
) -> list[GoldenResult]:
    """Run regression tests against a golden dataset."""
    results: list[GoldenResult] = []
    for example in examples:
        response = await agent.run(example.prompt, session_id=example.session_id)
        actual = response.content
        if example.mode == "contains":
            passed = example.expected in actual
        else:
            passed = actual == example.expected
        results.append(
            GoldenResult(
                name=example.name,
                passed=passed,
                actual=actual,
                expected=example.expected,
                mode=example.mode,
            )
        )
    return results


def load_golden_dataset(path: str | Path) -> list[GoldenExample]:
    """Load golden examples from JSON."""
    data = json.loads(Path(path).read_text())
    return [GoldenExample(**item) for item in data]


@dataclass
class JudgeCase:
    """Single evaluation case for LLM-as-judge."""

    name: str
    prompt: str
    criteria: str
    reference: str | None = None
    session_id: str | None = None


@dataclass
class JudgeResult:
    """LLM-as-judge evaluation result."""

    name: str
    score: float
    passed: bool
    rationale: str
    answer: str


class LLMJudge:
    """Evaluate agent answers using another agent or provider."""

    def __init__(self, judge: Agent | Callable[[str], Any], *, pass_threshold: float = 0.7) -> None:
        self.judge = judge
        self.pass_threshold = pass_threshold

    async def evaluate(self, *, prompt: str, answer: str, criteria: str, reference: str | None = None) -> tuple[float, str]:
        if isinstance(self.judge, Agent):
            judge_prompt = (
                "You are an impartial evaluator. Return strict JSON with keys "
                '"score" (0..1) and "rationale".\n'
                f"Prompt: {prompt}\n"
                f"Answer: {answer}\n"
                f"Criteria: {criteria}\n"
                f"Reference: {reference or ''}\n"
            )
            response = await self.judge.run(judge_prompt)
            return self._parse_judgement(response.content)
        verdict = self.judge(answer)
        if isinstance(verdict, tuple) and len(verdict) == 2:
            return float(verdict[0]), str(verdict[1])
        return float(verdict), "custom_judge"

    def _parse_judgement(self, content: str) -> tuple[float, str]:
        try:
            payload = json.loads(content)
            score = float(payload.get("score", 0))
            rationale = str(payload.get("rationale", ""))
            return max(0.0, min(1.0, score)), rationale
        except Exception:
            lowered = content.lower()
            if "excellent" in lowered or "pass" in lowered:
                return 0.9, content
            if "fail" in lowered or "incorrect" in lowered:
                return 0.1, content
            return 0.5, content


async def run_llm_judge_eval(
    agent: Agent,
    cases: list[JudgeCase],
    judge: LLMJudge,
) -> list[JudgeResult]:
    """Evaluate answers with an LLM-as-judge workflow."""
    results: list[JudgeResult] = []
    for case in cases:
        answer = await agent.run(case.prompt, session_id=case.session_id)
        score, rationale = await judge.evaluate(
            prompt=case.prompt,
            answer=answer.content,
            criteria=case.criteria,
            reference=case.reference,
        )
        results.append(
            JudgeResult(
                name=case.name,
                score=score,
                passed=score >= judge.pass_threshold,
                rationale=rationale,
                answer=answer.content,
            )
        )
    return results


class IsolatedTestEnvironment(AbstractContextManager["IsolatedTestEnvironment"]):
    """Sandboxed environment for safe agent evaluation."""

    def __init__(self, *, name: str = "logos-adk-sandbox") -> None:
        self.name = name
        self.root = Path(tempfile.mkdtemp(prefix=f"{name}-"))
        self.workspace = self.root / "workspace"
        self.workspace.mkdir(parents=True, exist_ok=True)
        self.state_dir = self.root / "state"
        self.state_dir.mkdir(parents=True, exist_ok=True)
        self._original_cwd = Path.cwd()
        self._original_env = os.environ.copy()

    def __enter__(self) -> IsolatedTestEnvironment:
        os.chdir(self.workspace)
        os.environ["LOGOS_ADK_SANDBOX"] = "1"
        os.environ["LOGOS_ADK_STATE_PATH"] = str(self.state_dir / "state.db")
        os.environ.pop("LOGOS_ADK_DATABASE_URL", None)
        os.environ.pop("DATABASE_URL", None)
        _default_runtime_holder.reset()
        reset_observability_store().reset()
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        os.chdir(self._original_cwd)
        os.environ.clear()
        os.environ.update(self._original_env)
        _default_runtime_holder.reset()
        reset_observability_store().reset()
        shutil.rmtree(self.root, ignore_errors=True)
        return None


__all__ = [
    "GoldenExample",
    "GoldenResult",
    "IntegrationResult",
    "IntegrationScenario",
    "IntegrationStep",
    "IsolatedTestEnvironment",
    "JudgeCase",
    "JudgeResult",
    "LLMJudge",
    "LoadTestConfig",
    "LoadTestReport",
    "load_golden_dataset",
    "run_golden_dataset",
    "run_integration_suite",
    "run_llm_judge_eval",
    "run_load_test",
]
