from pydantic import ValidationError

from dual_agent.domain import FixArtifact, ImplementationArtifact, PlanArtifact, ReviewArtifact

from .base import AgentRequest, AgentResult, Reply


SCHEMAS = {
    "PLAN": PlanArtifact,
    "PLAN_REVIEW": ReviewArtifact,
    "IMPLEMENT": ImplementationArtifact,
    "FIX": FixArtifact,
    "REVIEW": ReviewArtifact,
    "FINAL_VERIFY": ReviewArtifact,
}


class FakeAgentAdapter:
    def __init__(self, responses: dict[str, str], answer: str | None = None) -> None:
        self.responses = responses
        self.answer = answer or "（假适配器的回答）收到，我会按这个方向准备方案。"
        self.asked: list[str] = []

    def ask(self, prompt: str, cwd, session_id=None, on_start=None) -> Reply:
        self.asked.append(prompt)
        return Reply(self.answer)

    def run(self, request: AgentRequest, on_start=None) -> AgentResult:
        try:
            raw = self.responses[request.stage]
            artifact = SCHEMAS[request.stage].model_validate_json(raw)
            return AgentResult("SUCCESS", artifact=artifact)
        except (KeyError, ValidationError, ValueError) as error:
            return AgentResult("PROTOCOL", error=str(error))