"""Unit tests for the safe-genai pipeline."""

from __future__ import annotations

from unittest.mock import MagicMock

import pytest
from pydantic import BaseModel

from safe_genai.exceptions import (
    SafeGenAIFormatterError,
    SafeGenAIProviderError,
    SafeGenAIValidationError,
)
from safe_genai.pipeline.formatter import Formatter
from safe_genai.pipeline.generator import Generator
from safe_genai.pipeline.orchestrator import Pipeline
from safe_genai.pipeline.validator import Validator, _ValidationResult
from safe_genai.providers.base import BaseProvider
from safe_genai.result import AgentResult, AgentUsage, Usage


class Answer(BaseModel):
    value: str


def _usage(inp: int = 10, out: int = 5) -> Usage:
    return Usage(input_tokens=inp, output_tokens=out)


class FakeProvider(BaseProvider):
    """Configurable test double for BaseProvider."""

    def __init__(
        self,
        text_response: str = "raw text",
        structured_response: BaseModel | None = None,
        raise_on_text: Exception | None = None,
        raise_on_structured: Exception | None = None,
        text_usage: Usage | None = None,
        structured_usage: Usage | None = None,
    ) -> None:
        self._text_response = text_response
        self._structured_response = structured_response
        self._raise_on_text = raise_on_text
        self._raise_on_structured = raise_on_structured
        self._text_usage = text_usage or _usage(10, 5)
        self._structured_usage = structured_usage or _usage(20, 8)

    def generate_text(
        self, model, system_prompt, user_prompt, *, enable_grounding=False
    ) -> tuple[str, Usage]:
        if self._raise_on_text:
            raise self._raise_on_text
        return self._text_response, self._text_usage

    def generate_structured(
        self, model, system_prompt, user_prompt, output_schema
    ) -> tuple[BaseModel, Usage]:
        if self._raise_on_structured:
            raise self._raise_on_structured
        return self._structured_response or output_schema(value="ok"), self._structured_usage

    async def agenerate_text(
        self, model, system_prompt, user_prompt, *, enable_grounding=False
    ) -> tuple[str, Usage]:
        return self.generate_text(
            model, system_prompt, user_prompt, enable_grounding=enable_grounding
        )

    async def agenerate_structured(
        self, model, system_prompt, user_prompt, output_schema
    ) -> tuple[BaseModel, Usage]:
        return self.generate_structured(model, system_prompt, user_prompt, output_schema)


# ---------------------------------------------------------------------------
# Generator
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_generator_returns_text():
    provider = FakeProvider(text_response="hello world")
    gen = Generator(provider, "model-x", "sys")
    text, usage = await gen.agenerate("prompt")
    assert text == "hello world"
    assert isinstance(usage, Usage)


# ---------------------------------------------------------------------------
# Formatter
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_formatter_returns_model():
    expected = Answer(value="42")
    provider = FakeProvider(structured_response=expected)
    fmt = Formatter(provider, "model-x", Answer)
    result, usage = await fmt.aformat("raw text")
    assert result == expected
    assert isinstance(usage, Usage)


@pytest.mark.asyncio
async def test_formatter_raises_formatter_error_on_provider_failure():
    provider = FakeProvider(raise_on_structured=RuntimeError("boom"))
    fmt = Formatter(provider, "model-x", Answer)
    with pytest.raises(SafeGenAIFormatterError) as exc_info:
        await fmt.aformat("raw text")
    assert exc_info.value.schema_name == "Answer"
    assert "boom" in str(exc_info.value.cause)


# ---------------------------------------------------------------------------
# Validator
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_validator_passes():
    passing = _ValidationResult(passed=True, feedback="looks good")
    provider = FakeProvider(structured_response=passing)
    val = Validator(provider, "model-x")
    result, usage = await val.avalidate(Answer(value="42"))
    assert result.passed is True
    assert result.feedback == "looks good"
    assert isinstance(usage, Usage)


@pytest.mark.asyncio
async def test_validator_fails():
    failing = _ValidationResult(passed=False, feedback="bad output")
    provider = FakeProvider(structured_response=failing)
    val = Validator(provider, "model-x")
    result, _ = await val.avalidate(Answer(value="wrong"))
    assert result.passed is False


# ---------------------------------------------------------------------------
# Pipeline retry logic
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_pipeline_succeeds_first_try():
    passing = _ValidationResult(passed=True, feedback="great")
    formatted = Answer(value="4")

    # Two calls to generate_structured: formatter then validator
    call_count = 0

    class SequentialProvider(FakeProvider):
        def generate_structured(self, model, system_prompt, user_prompt, output_schema):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return formatted, _usage(20, 8)
            return passing, _usage(15, 6)

    provider = SequentialProvider(text_response="raw", text_usage=_usage(10, 5))
    gen = Generator(provider, "m", "sys")
    fmt = Formatter(provider, "m", Answer)
    val = Validator(provider, "m")
    pipeline = Pipeline(gen, fmt, val, max_retries=0)

    result = await pipeline.aexecute("2+2")
    assert isinstance(result, AgentResult)
    assert result.attempts == 1
    assert result.output == formatted
    assert result.validation_feedback == "great"
    assert result.raw_response == "raw"
    # usage breakdown
    assert result.usage.generator == Usage(10, 5)
    assert result.usage.formatter == Usage(20, 8)
    assert result.usage.validator == Usage(15, 6)
    assert result.usage.total == Usage(45, 19)


@pytest.mark.asyncio
async def test_pipeline_retries_and_succeeds():
    failing = _ValidationResult(passed=False, feedback="no good")
    passing = _ValidationResult(passed=True, feedback="ok now")
    formatted = Answer(value="4")

    call_count = 0

    class RetryProvider(FakeProvider):
        def generate_structured(self, model, system_prompt, user_prompt, output_schema):
            nonlocal call_count
            call_count += 1
            # Attempts: fmt(1), val(2) fail; fmt(3), val(4) pass
            if call_count in (1, 3):
                return formatted, _usage(20, 8)
            elif call_count == 2:
                return failing, _usage(15, 6)
            else:
                return passing, _usage(15, 6)

    provider = RetryProvider(text_response="raw", text_usage=_usage(10, 5))
    gen = Generator(provider, "m", "sys")
    fmt = Formatter(provider, "m", Answer)
    val = Validator(provider, "m")
    pipeline = Pipeline(gen, fmt, val, max_retries=1)

    result = await pipeline.aexecute("prompt")
    assert result.attempts == 2
    assert result.validation_feedback == "ok now"
    # usage accumulates across both attempts
    assert result.usage.generator == Usage(20, 10)  # 2 generator calls
    assert result.usage.formatter == Usage(40, 16)  # 2 formatter calls
    assert result.usage.validator == Usage(30, 12)  # 2 validator calls


@pytest.mark.asyncio
async def test_pipeline_raises_validation_error_after_all_retries():
    failing = _ValidationResult(passed=False, feedback="still bad")
    formatted = Answer(value="wrong")

    call_count = 0

    class AlwaysFailProvider(FakeProvider):
        def generate_structured(self, model, system_prompt, user_prompt, output_schema):
            nonlocal call_count
            call_count += 1
            if call_count % 2 == 1:
                return formatted, _usage()
            return failing, _usage()

    provider = AlwaysFailProvider(text_response="raw")
    gen = Generator(provider, "m", "sys")
    fmt = Formatter(provider, "m", Answer)
    val = Validator(provider, "m")
    pipeline = Pipeline(gen, fmt, val, max_retries=2)

    with pytest.raises(SafeGenAIValidationError) as exc_info:
        await pipeline.aexecute("prompt")
    err = exc_info.value
    assert err.attempts == 3
    assert err.feedback == "still bad"
    assert err.output == formatted
    assert err.raw_response == "raw"


@pytest.mark.asyncio
async def test_pipeline_formatter_error_not_retried():
    """SafeGenAIFormatterError must propagate immediately without retry."""
    call_count = 0

    class FormatterBombProvider(FakeProvider):
        def generate_structured(self, model, system_prompt, user_prompt, output_schema):
            nonlocal call_count
            call_count += 1
            raise RuntimeError("formatter exploded")

    provider = FormatterBombProvider(text_response="raw")
    gen = Generator(provider, "m", "sys")
    fmt = Formatter(provider, "m", Answer)
    val = Validator(provider, "m")
    pipeline = Pipeline(gen, fmt, val, max_retries=5)

    with pytest.raises(SafeGenAIFormatterError):
        await pipeline.aexecute("prompt")
    # Only one structured call made (formatter on attempt 1)
    assert call_count == 1


# ---------------------------------------------------------------------------
# Usage dataclass
# ---------------------------------------------------------------------------


def test_usage_addition():
    a = Usage(input_tokens=10, output_tokens=5)
    b = Usage(input_tokens=20, output_tokens=8)
    total = a + b
    assert total.input_tokens == 30
    assert total.output_tokens == 13


def test_agent_usage_total():
    usage = AgentUsage(
        generator=Usage(100, 50),
        formatter=Usage(200, 80),
        validator=Usage(150, 60),
    )
    assert usage.total == Usage(450, 190)


# ---------------------------------------------------------------------------
# Exception attributes
# ---------------------------------------------------------------------------


def test_validation_error_attrs():
    model = Answer(value="x")
    err = SafeGenAIValidationError(output=model, feedback="bad", attempts=3, raw_response="raw")
    assert err.output is model
    assert err.feedback == "bad"
    assert err.attempts == 3
    assert err.raw_response == "raw"
    assert "3 attempt" in str(err)


def test_formatter_error_attrs():
    cause = ValueError("parse fail")
    err = SafeGenAIFormatterError(raw_response="text", schema_name="MySchema", cause=cause)
    assert err.raw_response == "text"
    assert err.schema_name == "MySchema"
    assert err.cause is cause


def test_provider_error_attrs():
    cause = ConnectionError("timeout")
    err = SafeGenAIProviderError(provider_name="GoogleGenAI", cause=cause)
    assert err.provider_name == "GoogleGenAI"
    assert err.cause is cause


# ---------------------------------------------------------------------------
# Provider abstraction
# ---------------------------------------------------------------------------


def test_base_provider_is_abstract():
    """BaseProvider cannot be instantiated directly."""
    with pytest.raises(TypeError):
        BaseProvider()  # type: ignore[abstract]


# ---------------------------------------------------------------------------
# GoogleGenAIProvider init modes
# ---------------------------------------------------------------------------


def test_google_provider_requires_api_key_or_env(monkeypatch):
    """Should raise ValueError when no api_key and GEMINI_API_KEY is unset."""
    monkeypatch.delenv("GEMINI_API_KEY", raising=False)
    from safe_genai.providers.google import GoogleGenAIProvider

    with pytest.raises(ValueError, match="GEMINI_API_KEY"):
        GoogleGenAIProvider()


def test_google_provider_vertexai_requires_project(monkeypatch):
    """Should raise ValueError when vertexai=True and no project is provided."""
    monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False)
    from safe_genai.providers.google import GoogleGenAIProvider

    with pytest.raises(ValueError, match="vertex_project"):
        GoogleGenAIProvider(vertexai=True)


def test_google_provider_vertexai_accepts_env_project(monkeypatch):
    """Should not raise when GOOGLE_CLOUD_PROJECT is set and vertexai=True."""
    from unittest.mock import patch

    monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "my-project")
    mock_client = MagicMock()
    with patch("google.genai.Client", return_value=mock_client) as mock_ctor:
        from safe_genai.providers.google import GoogleGenAIProvider

        GoogleGenAIProvider(vertexai=True)
        mock_ctor.assert_called_once_with(
            vertexai=True, project="my-project", location="us-central1"
        )
