from __future__ import annotations

from typing import Generic, TypeVar

from safe_genai.exceptions import SafeGenAIValidationError
from safe_genai.pipeline.formatter import Formatter
from safe_genai.pipeline.generator import Generator
from safe_genai.pipeline.validator import Validator
from safe_genai.result import AgentResult, AgentUsage, Usage

T = TypeVar("T")


class Pipeline(Generic[T]):
    def __init__(
        self,
        generator: Generator,
        formatter: Formatter,
        validator: Validator,
        max_retries: int = 0,
    ) -> None:
        self._generator = generator
        self._formatter = formatter
        self._validator = validator
        self._max_retries = max_retries

    def execute(self, user_prompt: str) -> AgentResult[T]:
        formatted = None
        result = None
        raw = ""
        gen_usage = Usage()
        fmt_usage = Usage()
        val_usage = Usage()
        for attempt in range(1, self._max_retries + 2):
            raw, u = self._generator.generate(user_prompt)
            gen_usage = gen_usage + u
            formatted, u = self._formatter.format(
                raw
            )  # raises SafeGenAIFormatterError (not retried)
            fmt_usage = fmt_usage + u
            result, u = self._validator.validate(formatted)
            val_usage = val_usage + u
            if result.passed:
                return AgentResult(
                    output=formatted,  # type: ignore[arg-type]
                    raw_response=raw,
                    validation_feedback=result.feedback,
                    attempts=attempt,
                    usage=AgentUsage(generator=gen_usage, formatter=fmt_usage, validator=val_usage),
                )
        raise SafeGenAIValidationError(
            output=formatted,  # type: ignore[arg-type]
            feedback=result.feedback if result else "",
            attempts=self._max_retries + 1,
            raw_response=raw,
        )

    async def aexecute(self, user_prompt: str) -> AgentResult[T]:
        formatted = None
        result = None
        raw = ""
        gen_usage = Usage()
        fmt_usage = Usage()
        val_usage = Usage()
        for attempt in range(1, self._max_retries + 2):
            raw, u = await self._generator.agenerate(user_prompt)
            gen_usage = gen_usage + u
            formatted, u = await self._formatter.aformat(
                raw
            )  # raises SafeGenAIFormatterError (not retried)
            fmt_usage = fmt_usage + u
            result, u = await self._validator.avalidate(formatted)
            val_usage = val_usage + u
            if result.passed:
                return AgentResult(
                    output=formatted,  # type: ignore[arg-type]
                    raw_response=raw,
                    validation_feedback=result.feedback,
                    attempts=attempt,
                    usage=AgentUsage(generator=gen_usage, formatter=fmt_usage, validator=val_usage),
                )
        raise SafeGenAIValidationError(
            output=formatted,  # type: ignore[arg-type]
            feedback=result.feedback if result else "",
            attempts=self._max_retries + 1,
            raw_response=raw,
        )
