from __future__ import annotations

import asyncio
from typing import Generic, TypeVar

from safe_genai.config import AgentConfig
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
from safe_genai.providers.base import BaseProvider
from safe_genai.providers.google import GoogleGenAIProvider
from safe_genai.result import AgentResult

T = TypeVar("T")


class SafeAgent(Generic[T]):
    def __init__(
        self,
        output_schema: type[T],
        system_prompt: str,
        additional_formatter_instructions: str | None = None,
        additional_validator_instructions: str | None = None,
        max_retries: int = 0,
        config: AgentConfig | None = None,
        provider: BaseProvider | None = None,
    ) -> None:
        cfg = config or AgentConfig()

        resolved_provider = provider or GoogleGenAIProvider(
            api_key=cfg.api_key,
            vertexai=cfg.vertexai,
            vertex_project=cfg.vertex_project,
            vertex_location=cfg.vertex_location,
        )

        generator = Generator(
            provider=resolved_provider,
            model=cfg.generator_model,
            system_prompt=system_prompt,
        )
        formatter = Formatter(
            provider=resolved_provider,
            model=cfg.formatter_model,
            output_schema=output_schema,  # type: ignore[arg-type]
            additional_instructions=additional_formatter_instructions,
        )
        validator = Validator(
            provider=resolved_provider,
            model=cfg.validator_model,
            additional_instructions=additional_validator_instructions,
        )

        self._pipeline: Pipeline[T] = Pipeline(
            generator=generator,
            formatter=formatter,
            validator=validator,
            max_retries=max_retries,
        )

    def run(self, user_prompt: str) -> AgentResult[T]:
        """Synchronous pipeline run. Raises SafeGenAIValidationError on failure."""
        return asyncio.run(self.run_async(user_prompt))

    async def run_async(self, user_prompt: str) -> AgentResult[T]:
        """Async pipeline run. Same contract as run()."""
        return await self._pipeline.aexecute(user_prompt)
