"""
Loom message schemas — the canonical wire format.

All inter-actor communication is typed through these Pydantic models.
Actors ONLY communicate through these message types; raw dicts or
ad-hoc JSON are forbidden. This enforces a contract-driven architecture
where every message is validatable at compile time.

Message flow:
    Client/CLI  ──OrchestratorGoal──>  Orchestrator
    Orchestrator  ──TaskMessage──>  Router  ──TaskMessage──>  Worker
    Worker  ──TaskResult──>  Orchestrator

See Also:
    loom.core.contracts — JSON Schema validation for payload/output dicts
    loom.bus.nats_adapter — NATS subject conventions for message routing
"""

from __future__ import annotations

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

from pydantic import BaseModel, Field


class TaskPriority(StrEnum):
    """Priority levels for task scheduling (not yet enforced by router)."""

    LOW = "low"
    NORMAL = "normal"
    HIGH = "high"
    CRITICAL = "critical"


class TaskStatus(StrEnum):
    """Lifecycle states for a task.

    State transitions:
        PENDING -> PROCESSING -> COMPLETED
                               -> FAILED
                               -> RETRY -> PROCESSING (not yet implemented)
    """

    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    RETRY = "retry"  # Transition: FAILED -> RETRY -> PROCESSING (handled by TaskWorker)


class ModelTier(StrEnum):
    """Which model tier should handle this task.

    Tiers map to backend instances configured at startup:
        LOCAL    -> OllamaBackend (e.g., llama3.2:3b)
        STANDARD -> AnthropicBackend (e.g., Claude Sonnet)
        FRONTIER -> AnthropicBackend (e.g., Claude Opus)

    The router may override the tier via tier_overrides in router_rules.yaml.
    """

    LOCAL = "local"  # Small local model (Ollama, llama.cpp)
    STANDARD = "standard"  # Mid-tier API model
    FRONTIER = "frontier"  # Top-tier model (Claude Opus, GPT-4, etc.)


class TaskMessage(BaseModel):
    """Message sent TO a worker actor.

    The payload dict must conform to the worker's input_schema (JSON Schema).
    Contract validation happens in TaskWorker.handle_message(), not here.
    """

    task_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    parent_task_id: str | None = None  # Links subtask to orchestrator's goal
    worker_type: str  # Which worker config to use (e.g., "summarizer", "doc_extractor")
    payload: dict[str, Any]  # Structured input — must match worker's input_schema
    model_tier: ModelTier = ModelTier.STANDARD
    priority: TaskPriority = TaskPriority.NORMAL
    max_retries: int = 2  # Max retry attempts before permanent failure
    retry_count: int = 0  # Incremented on each retry by TaskWorker
    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
    request_id: str | None = None  # Correlates all tasks from the same goal (set by pipeline)
    metadata: dict[str, Any] = Field(default_factory=dict)  # Routing hints, pipeline context


class TaskResult(BaseModel):
    """Message sent FROM a worker actor after processing.

    Published to: loom.results.{parent_task_id or 'default'}
    The output dict must conform to the worker's output_schema (JSON Schema).
    """

    task_id: str
    parent_task_id: str | None = None
    worker_type: str
    status: TaskStatus
    output: dict[str, Any] | None = None  # Structured output — must match worker's output_schema
    error: str | None = None  # Human-readable error message on failure
    model_used: str | None = None  # Actual model that processed this (e.g., "llama3.2:3b")
    token_usage: dict[str, int] = Field(
        default_factory=dict
    )  # {"prompt_tokens": N, "completion_tokens": N}
    processing_time_ms: int = 0
    completed_at: datetime = Field(default_factory=lambda: datetime.now(UTC))


class OrchestratorGoal(BaseModel):
    """Top-level goal submitted to an orchestrator.

    Published to: loom.goals.incoming
    The orchestrator (PipelineOrchestrator or OrchestratorActor) picks this up,
    decomposes it into TaskMessages, and synthesizes results.

    The context dict carries domain-specific data (e.g., file_ref for doc processing).
    """

    goal_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    instruction: str  # Natural language goal
    context: dict[str, Any] = Field(
        default_factory=dict
    )  # Domain data (file_ref, categories, etc.)
    request_id: str | None = None  # Optional correlation ID for tracing goal→task chains
    priority: TaskPriority = TaskPriority.NORMAL
    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))


class CheckpointState(BaseModel):
    """Compressed orchestrator state for self-summarization.

    When the orchestrator's conversation history exceeds a token threshold,
    CheckpointManager compresses it into this structure and persists to Valkey.
    The orchestrator can then "reboot" with a fresh context containing only
    the checkpoint + a small recent-interactions window.

    See: loom.orchestrator.checkpoint.CheckpointManager
    """

    goal_id: str
    original_instruction: str
    executive_summary: str  # High-level status (always short)
    completed_tasks: list[dict[str, Any]]  # Key outcomes only, not full results
    pending_tasks: list[dict[str, Any]]  # What remains
    open_issues: list[str]  # Conflicts, blockers, uncertainties
    decisions_made: list[str]  # Important choices and rationale
    context_token_count: int  # Tokens at time of checkpoint
    checkpoint_number: int
    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
