from __future__ import annotations

import json
from typing import Any

from agents.models import Observation, ReflectionAction, ReflectionOutcome, VerificationVerdict
from runtime.executor import StepResult
from runtime.planner import PlannedStep
from semantic.llm_client import AnthropicSemanticClient, EnrichmentUnavailable

REFLECTOR_SYSTEM_PROMPT = """You are a reflection agent. A step failed to achieve its intended effect despite mechanical success. Your job is to propose either:
1. A distinct recovery action (different from the failed step) that might fix the issue
2. A replan request if no bounded corrective action is identifiable

You must NEVER simply re-issue the identical failed step.

Respond with JSON: {"replan": bool, "action": null or {"action": "...", "description": "...", "args": {...}}, "reasoning": "explanation"}"""

REFLECTOR_JSON_SCHEMA: dict[str, Any] = {
    "type": "object",
    "additionalProperties": False,
    "required": ["replan", "reasoning"],
    "properties": {
        "replan": {"type": "boolean"},
        "action": {
            "type": ["object", "null"],
            "properties": {
                "action": {"type": "string"},
                "description": {"type": "string"},
                "args": {"type": "object"},
            },
        },
        "reasoning": {"type": "string", "minLength": 1},
    },
}


class Reflector:
    def __init__(self, *, client: AnthropicSemanticClient | None = None) -> None:
        self._client = client

    def reflect(
        self,
        *,
        observation: Observation,
        planned_step: PlannedStep,
        step_result: StepResult,
        verdict: VerificationVerdict,
    ) -> ReflectionOutcome:
        if self._client is None:
            return self._fallback_reflect(
                observation=observation,
                planned_step=planned_step,
                step_result=step_result,
                verdict=verdict,
            )

        try:
            return self._llm_reflect(
                observation=observation,
                planned_step=planned_step,
                step_result=step_result,
                verdict=verdict,
            )
        except EnrichmentUnavailable:
            return self._fallback_reflect(
                observation=observation,
                planned_step=planned_step,
                step_result=step_result,
                verdict=verdict,
            )

    def _llm_reflect(
        self,
        *,
        observation: Observation,
        planned_step: PlannedStep,
        step_result: StepResult,
        verdict: VerificationVerdict,
    ) -> ReflectionOutcome:
        user_prompt = json.dumps(
            {
                "observation": observation.to_dict(),
                "planned_step": {
                    "action": planned_step.action,
                    "description": planned_step.description,
                    "args": planned_step.args,
                    "expected_text": planned_step.expected_text,
                },
                "step_result": step_result.to_dict(),
                "verdict": verdict.to_dict(),
            },
            ensure_ascii=False,
        )
        response = self._client._create_message(  # type: ignore[union-attr]
            {"prompt": user_prompt, "system": REFLECTOR_SYSTEM_PROMPT},
            timeout=10.0,
        )
        payload = _extract_outcome(response)
        outcome = ReflectionOutcome.from_dict(payload)
        if not outcome.replan and outcome.action is not None:
            if outcome.action.action == planned_step.action and outcome.action.args == planned_step.args:
                return ReflectionOutcome(
                    replan=True,
                    reasoning="Reflector proposed identical action; requesting replan instead.",
                )
        return outcome

    def _fallback_reflect(
        self,
        *,
        observation: Observation,
        planned_step: PlannedStep,
        step_result: StepResult,
        verdict: VerificationVerdict,
    ) -> ReflectionOutcome:
        return ReflectionOutcome(
            replan=True,
            reasoning=f"Step '{planned_step.description}' failed semantically; replanning (fallback).",
        )


def _extract_outcome(response: Any) -> dict[str, Any]:
    if isinstance(response, dict) and "replan" in response:
        return response
    content = getattr(response, "content", None)
    if isinstance(content, list):
        for block in content:
            text = getattr(block, "text", None)
            if isinstance(text, str):
                try:
                    decoded = json.loads(text)
                    if isinstance(decoded, dict) and "replan" in decoded:
                        return decoded
                except json.JSONDecodeError:
                    continue
            if isinstance(block, dict) and "replan" in block:
                return block
    if isinstance(content, str):
        try:
            decoded = json.loads(content)
            if isinstance(decoded, dict) and "replan" in decoded:
                return decoded
        except json.JSONDecodeError:
            pass
    raise ValueError("could not extract reflection outcome from LLM response")