"""Critic-reviser loop for answer improvement."""
from __future__ import annotations

import asyncio
import json
from typing import TYPE_CHECKING, Any, Awaitable, Callable, Union

if TYPE_CHECKING:
    from logos_adk.agent import Agent

from .models import CriticIssue, CriticReview, CriticReviserResult, RevisionSuggestion


ActorType = Union["Agent", Callable[[str], Any], Callable[[str], Awaitable[Any]]]


class CriticReviserLoop:
    """Critic detects defects; reviser proposes an improved answer/prompt."""

    def __init__(
        self,
        critic: ActorType,
        reviser: ActorType,
    ) -> None:
        """Initialize critic-reviser loop.
        
        Args:
            critic: Agent or function that reviews answers
            reviser: Agent or function that proposes improvements
        """
        self.critic = critic
        self.reviser = reviser

    async def improve(
        self,
        *,
        prompt: str,
        answer: str,
        reference: str | None = None,
        failure_types: list[str] | None = None,
    ) -> CriticReviserResult:
        """Run full critic-reviser improvement cycle."""
        review = await self.critique(
            prompt=prompt,
            answer=answer,
            reference=reference,
            failure_types=failure_types,
        )
        if not review.should_revise:
            return CriticReviserResult(review=review, revision=RevisionSuggestion(), revised=False)
        revision = await self.revise(
            prompt=prompt,
            answer=answer,
            review=review,
            reference=reference,
        )
        return CriticReviserResult(review=review, revision=revision, revised=bool(revision.improved_answer or revision.prompt_patch))

    async def critique(
        self,
        *,
        prompt: str,
        answer: str,
        reference: str | None = None,
        failure_types: list[str] | None = None,
    ) -> CriticReview:
        """Critique an answer for defects."""
        critic_prompt = (
            "Review the answer and return JSON with keys "
            '"should_revise", "summary", "issues", and optional "prompt_patch".\n'
            f"Prompt: {prompt}\nAnswer: {answer}\nReference: {reference or ''}\n"
            f"Failure types: {', '.join(failure_types or [])}\n"
        )
        payload = await self._invoke(self.critic, critic_prompt)
        return self._parse_review(payload)

    async def revise(
        self,
        *,
        prompt: str,
        answer: str,
        review: CriticReview,
        reference: str | None = None,
    ) -> RevisionSuggestion:
        """Revise an answer based on critic review."""
        reviser_prompt = (
            "Revise the answer using the critic review and return JSON with keys "
            '"improved_answer", optional "prompt_patch", and optional "rationale".\n'
            f"Prompt: {prompt}\nAnswer: {answer}\n"
            f"Critic summary: {review.summary}\n"
            f"Reference: {reference or ''}\n"
        )
        payload = await self._invoke(self.reviser, reviser_prompt)
        return self._parse_revision(payload)

    async def _invoke(self, actor: ActorType, prompt: str) -> Any:
        """Invoke actor (agent or function) with prompt."""
        # Import here to avoid circular import
        from logos_adk.agent import Agent
        
        if isinstance(actor, Agent):
            response = await actor.run(prompt)
            return response.content
        result = actor(prompt)
        if asyncio.iscoroutine(result):
            result = await result
        return result

    def _parse_review(self, payload: Any) -> CriticReview:
        """Parse critic review from payload."""
        if isinstance(payload, CriticReview):
            return payload
        if isinstance(payload, dict):
            issues = [
                CriticIssue(
                    kind=str(item.get("kind", "issue")),
                    severity=str(item.get("severity", "medium")),
                    description=str(item.get("description", "")),
                )
                for item in (payload.get("issues") or [])
                if isinstance(item, dict)
            ]
            return CriticReview(
                should_revise=bool(payload.get("should_revise", bool(issues))),
                summary=str(payload.get("summary", "")),
                issues=issues,
                prompt_patch=payload.get("prompt_patch"),
            )
        if isinstance(payload, str):
            try:
                return self._parse_review(json.loads(payload))
            except Exception:
                issue = CriticIssue(kind="issue", severity="medium", description=payload)
                return CriticReview(should_revise=True, summary=payload, issues=[issue])
        raise ValueError(f"Unsupported critic payload: {type(payload)}")

    def _parse_revision(self, payload: Any) -> RevisionSuggestion:
        """Parse revision suggestion from payload."""
        if isinstance(payload, RevisionSuggestion):
            return payload
        if isinstance(payload, dict):
            return RevisionSuggestion(
                improved_answer=payload.get("improved_answer"),
                prompt_patch=payload.get("prompt_patch"),
                rationale=payload.get("rationale"),
            )
        if isinstance(payload, str):
            try:
                return self._parse_revision(json.loads(payload))
            except Exception:
                return RevisionSuggestion(improved_answer=payload)
        raise ValueError(f"Unsupported reviser payload: {type(payload)}")


__all__ = ["CriticReviserLoop"]
