from __future__ import annotations

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

from pydantic import BaseModel, Field

from nexus.core.state.cold_layer import ColdLayer


class Skill(BaseModel):
    """Stored high-quality execution skill."""

    skill_id: str = Field(default_factory=lambda: str(uuid4()))
    agent_id: str
    task_type: str
    description: str
    execution_trace: str
    quality_score: float
    tags: list[str] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    use_count: int = 0


class SkillRegistry:
    """ColdLayer-backed skill persistence and retrieval."""

    def __init__(self, cold_layer: ColdLayer, quality_threshold: float = 0.8) -> None:
        self.cold_layer = cold_layer
        self.quality_threshold = quality_threshold

    async def store(
        self,
        agent_id: str,
        task_type: str,
        description: str,
        execution_trace: str,
        quality_score: float,
        tags: list[str] | None = None,
    ) -> Skill | None:
        """Persist high-quality skill; return None if below threshold."""
        if quality_score < self.quality_threshold:
            return None
        skill = Skill(
            agent_id=agent_id,
            task_type=task_type,
            description=description,
            execution_trace=execution_trace,
            quality_score=quality_score,
            tags=list(tags or []),
        )
        key = f"skill:{skill.skill_id}"
        self.cold_layer.store(
            key=key,
            content=skill.model_dump_json(),
            metadata={
                "agent_id": agent_id,
                "task_type": task_type,
                "tags": skill.tags,
                "quality_score": quality_score,
                "created_at": skill.created_at.isoformat(),
                "skill_id": skill.skill_id,
            },
        )
        return skill

    async def retrieve(self, query: str, n_results: int = 3, min_quality: float = 0.7) -> list[Skill]:
        """Retrieve skills by semantic search and quality filter."""
        rows = self.cold_layer.retrieve(query, n_results=n_results)
        out: list[Skill] = []
        for row in rows:
            meta = row.get("metadata", {})
            if float(meta.get("quality_score", 0.0)) < min_quality:
                continue
            try:
                out.append(Skill.model_validate_json(row.get("content", "{}")))
            except ValueError:
                continue
        return out

    async def to_few_shot_examples(self, query: str, n: int = 2) -> str:
        """Format top-n skills as few-shot examples."""
        skills = await self.retrieve(query, n_results=n)
        chunks: list[str] = []
        for i, skill in enumerate(skills, start=1):
            chunks.append(f"## Example {i}: {skill.description}\n{skill.execution_trace}\n")
        return "\n".join(chunks)

    async def increment_use_count(self, skill_id: str) -> None:
        """Increment use count for a stored skill."""
        key = f"skill:{skill_id}"
        rows = self.cold_layer.retrieve(skill_id, n_results=20, where={"skill_id": skill_id})
        if not rows:
            return
        content = rows[0].get("content", "{}")
        try:
            skill = Skill.model_validate_json(content)
        except ValueError:
            return
        skill.use_count += 1
        self.cold_layer.store(
            key=key,
            content=skill.model_dump_json(),
            metadata={
                "agent_id": skill.agent_id,
                "task_type": skill.task_type,
                "tags": skill.tags,
                "quality_score": skill.quality_score,
                "created_at": skill.created_at.isoformat(),
                "skill_id": skill.skill_id,
                "use_count": skill.use_count,
            },
        )
