from __future__ import annotations

import json
from datetime import datetime, timezone
from uuid import uuid4

from pydantic import BaseModel, Field

from nexus.protocols.providers.base import AbstractProvider, LLMRequest, Message


class ReflectionRecord(BaseModel):
    """Structured post-task reflection record."""

    reflection_id: str = Field(default_factory=lambda: str(uuid4()))
    agent_id: str
    node_id: str
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    task_description: str
    what_went_well: list[str]
    what_went_poorly: list[str]
    suggestions: list[str]
    quality_score: float
    token_efficiency: float


class PostTaskReflector:
    """Generates post-task reflection using provider call."""

    def __init__(self, provider: AbstractProvider) -> None:
        self.provider = provider

    async def reflect(
        self,
        agent_id: str,
        node_id: str,
        task_description: str,
        task_result: str,
        tokens_used: int,
        errors_encountered: list[str],
    ) -> ReflectionRecord:
        """Generate structured reflection and compute token efficiency."""
        prompt = self._build_reflection_prompt(
            task_description=task_description,
            task_result=task_result,
            tokens_used=tokens_used,
            errors_encountered=errors_encountered,
        )
        req = LLMRequest(messages=[Message(role="user", content=prompt)], temperature=0.2, max_tokens=800)
        resp = await self.provider.complete(req)
        data = _extract_json(resp.content)

        quality = float(data.get("quality_score", 0.5))
        quality = max(0.0, min(1.0, quality))
        denom = max(tokens_used / 1000.0, 0.001)
        token_eff = quality / denom

        return ReflectionRecord(
            agent_id=agent_id,
            node_id=node_id,
            task_description=task_description,
            what_went_well=[str(x) for x in data.get("what_went_well", [])],
            what_went_poorly=[str(x) for x in data.get("what_went_poorly", [])],
            suggestions=[str(x) for x in data.get("suggestions", [])],
            quality_score=quality,
            token_efficiency=token_eff,
        )

    def _build_reflection_prompt(
        self,
        task_description: str,
        task_result: str,
        tokens_used: int,
        errors_encountered: list[str],
    ) -> str:
        """Build strict JSON reflection prompt."""
        return (
            "Return ONLY JSON with keys what_went_well, what_went_poorly, suggestions, quality_score.\n"
            f"Task: {task_description}\n"
            f"Result: {task_result}\n"
            f"Tokens used: {tokens_used}\n"
            f"Errors: {errors_encountered}\n"
        )


def _extract_json(text: str) -> dict:
    import re

    m = re.search(r"```json\s*(.*?)\s*```", text, flags=re.DOTALL)
    raw = m.group(1) if m else text
    try:
        parsed = json.loads(raw)
        return parsed if isinstance(parsed, dict) else {}
    except ValueError:
        return {}
