"""Core type definitions shared across the Nexus framework."""

from __future__ import annotations

import uuid
from datetime import datetime
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# IDs
# ---------------------------------------------------------------------------

def new_id() -> str:
    """Generate a short, URL-safe unique id."""
    return uuid.uuid4().hex[:16]


# ---------------------------------------------------------------------------
# Agent capability vector
# ---------------------------------------------------------------------------

class AgentCapability(BaseModel):
    """
    Describes what an agent can do — used for discovery and adaptive matching.

    Capabilities are modeled as a *vector* rather than tags, so that the
    TopologyEngine can compute similarity scores for optimal task assignment.
    """

    coding: float = Field(default=0.0, ge=0.0, le=1.0)
    reasoning: float = Field(default=0.0, ge=0.0, le=1.0)
    writing: float = Field(default=0.0, ge=0.0, le=1.0)
    research: float = Field(default=0.0, ge=0.0, le=1.0)
    data_analysis: float = Field(default=0.0, ge=0.0, le=1.0)
    math: float = Field(default=0.0, ge=0.0, le=1.0)
    creativity: float = Field(default=0.0, ge=0.0, le=1.0)
    communication: float = Field(default=0.0, ge=0.0, le=1.0)

    # Extensible: users can add custom dimensions
    custom: dict[str, float] = Field(default_factory=dict)

    def to_vector(self) -> list[float]:
        """Flatten to a numeric vector for similarity computation."""
        base = [
            self.coding,
            self.reasoning,
            self.writing,
            self.research,
            self.data_analysis,
            self.math,
            self.creativity,
            self.communication,
        ]
        # Append custom dimensions in sorted key order for determinism
        base.extend(v for _, v in sorted(self.custom.items()))
        return base

    def cosine_similarity(self, other: AgentCapability) -> float:
        """Compute cosine similarity between two capability vectors."""
        a = self.to_vector()
        b = other.to_vector()
        # Pad shorter vector with zeros
        max_len = max(len(a), len(b))
        a.extend([0.0] * (max_len - len(a)))
        b.extend([0.0] * (max_len - len(b)))

        dot = sum(x * y for x, y in zip(a, b))
        norm_a = sum(x * x for x in a) ** 0.5
        norm_b = sum(x * x for x in b) ** 0.5
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)

    def is_zero(self) -> bool:
        """Check if all capability dimensions are zero (i.e. unset)."""
        return all(v == 0.0 for v in self.to_vector())


# ---------------------------------------------------------------------------
# Task constraints
# ---------------------------------------------------------------------------

class TaskConstraints(BaseModel):
    """Constraints that govern how a task should be executed."""

    max_budget_usd: float | None = Field(
        default=None,
        description="Maximum cost in USD for this task",
    )
    max_tokens: int | None = Field(
        default=None,
        description="Maximum total tokens (input + output) for this task",
    )
    max_latency_seconds: float | None = Field(
        default=None,
        description="Maximum wall-clock time in seconds",
    )
    min_quality_score: float = Field(
        default=0.0,
        ge=0.0,
        le=1.0,
        description="Minimum acceptable quality score (0.0 - 1.0)",
    )
    required_capabilities: AgentCapability | None = Field(
        default=None,
        description="Capability vector describing what this task needs",
    )


# ---------------------------------------------------------------------------
# Agent Card (A2A-compatible)
# ---------------------------------------------------------------------------

class AgentCard(BaseModel):
    """
    Self-describing manifest for an agent, inspired by the A2A protocol spec.

    Can be serialized to JSON and served at /.well-known/agent-card.json
    """

    agent_id: str = Field(default_factory=new_id)
    name: str
    description: str = ""
    version: str = "0.1.0"
    capabilities: AgentCapability = Field(default_factory=AgentCapability)
    input_content_types: list[str] = Field(default_factory=lambda: ["text/plain"])
    output_content_types: list[str] = Field(default_factory=lambda: ["text/plain"])
    skills: list[str] = Field(default_factory=list)
    tools: list[str] = Field(default_factory=list)
    endpoint: str | None = None
    auth_required: bool = False
    created_at: datetime = Field(default_factory=datetime.utcnow)
    metadata: dict[str, Any] = Field(default_factory=dict)


# ---------------------------------------------------------------------------
# Quality score record
# ---------------------------------------------------------------------------

class QualityRecord(BaseModel):
    """Records an agent's performance on a specific task for the feedback loop."""

    agent_id: str
    task_id: str
    task_type: str = "general"
    score: float = Field(ge=0.0, le=1.0)
    tokens_used: int = 0
    cost_usd: float = 0.0
    latency_seconds: float = 0.0
    timestamp: datetime = Field(default_factory=datetime.utcnow)


# ---------------------------------------------------------------------------
# Token usage tracking
# ---------------------------------------------------------------------------

class TokenUsage(BaseModel):
    """Tracks token consumption for a single LLM call."""

    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cost_usd: float = 0.0
    model: str = ""

    def __add__(self, other: TokenUsage) -> TokenUsage:
        return TokenUsage(
            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
            completion_tokens=self.completion_tokens + other.completion_tokens,
            total_tokens=self.total_tokens + other.total_tokens,
            cost_usd=self.cost_usd + other.cost_usd,
            model=self.model or other.model,
        )


# ---------------------------------------------------------------------------
# Content parts (multi-modal, A2A-compatible)
# ---------------------------------------------------------------------------

class ContentType(StrEnum):
    TEXT = "text"
    IMAGE = "image"
    FILE = "file"
    JSON = "json"


class ContentPart(BaseModel):
    """
    Minimal content unit, inspired by A2A's Part abstraction.
    Supports text, images, files, and structured JSON data.
    """

    type: ContentType = ContentType.TEXT
    text: str | None = None
    data: dict[str, Any] | None = None
    url: str | None = None
    mime_type: str = "text/plain"
    metadata: dict[str, Any] = Field(default_factory=dict)

    @classmethod
    def from_text(cls, text: str) -> ContentPart:
        return cls(type=ContentType.TEXT, text=text)

    @classmethod
    def from_json(cls, data: dict[str, Any]) -> ContentPart:
        return cls(type=ContentType.JSON, data=data, mime_type="application/json")
