"""Pydantic models for the licensing subsystem."""
from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field


# ── Enums ────────────────────────────────────────────────────────────────────


class LicenseTier(str, Enum):
    COMMUNITY = "community"
    PROFESSIONAL = "professional"
    ENTERPRISE = "enterprise"


# ── Feature Descriptions ─────────────────────────────────────────────────────
# Community = free trial / open-source teaser
#   - AMD GPU optimization (one-time probe only, no continuous re-optimization)
#   - Dashboard, job routing, API access (basic kubernetes-like orchestration)
#   - Single probe per device — no scheduled re-probing or online learning
#
# Professional = full continuous optimization for all hardware
#   - All community features PLUS:
#   - Continuous optimization (scheduled probes, JEPA retraining, online learning)
#   - Auto-tuning (CPU, memory, scheduler, network — all devices)
#   - NVIDIA GPU, CPU, FPGA, memory subsystem optimization
#   - Codex AI agents
#
# Enterprise = unlimited scale + priority support
#   - All professional features PLUS:
#   - Unlimited nodes, multi-GPU, credential vault, priority support
#
# On license expiry: revert to community (AMD GPU one-shot stays, continuous stops)
# Dashboard and job routing NEVER get disabled — it's free simple kubernetes.

TIER_DEFAULTS: Dict[LicenseTier, Dict[str, Any]] = {
    LicenseTier.COMMUNITY: {
        "max_nodes": 0,  # unlimited nodes — free kubernetes
        "features": [
            "dashboard",
            "api_access",
            "job_routing",           # Free — basic job orchestration
            "amd_gpu_optimization",  # Free — one-time probe per AMD GPU
            "one_time_probe",        # Single probe per device, no continuous
            "credential_vault",      # Free — security should never be paywalled
        ],
    },
    LicenseTier.PROFESSIONAL: {
        "max_nodes": 50,
        "features": [
            "dashboard",
            "api_access",
            "job_routing",
            "amd_gpu_optimization",
            "one_time_probe",
            "credential_vault",
            "continuous_optimization",  # Scheduled re-probing + online learning
            "auto_tuning",              # Full system tuning (all hardware)
            "jepa_training",            # JEPA retraining + model sync
            "nvidia_optimization",      # NVIDIA GPU optimization
            "cpu_optimization",         # CPU/scheduler optimization
            "fpga_optimization",        # FPGA optimization
            "memory_optimization",      # Memory subsystem optimization
            # codex_agents is internal tooling, not customer-facing
        ],
    },
    LicenseTier.ENTERPRISE: {
        "max_nodes": 0,  # unlimited
        "features": [
            "dashboard",
            "api_access",
            "job_routing",
            "amd_gpu_optimization",
            "one_time_probe",
            "credential_vault",
            "continuous_optimization",
            "auto_tuning",
            "jepa_training",
            "nvidia_optimization",
            "cpu_optimization",
            "fpga_optimization",
            "memory_optimization",
            "codex_agents",
            "multi_gpu",                # Multi-GPU orchestration
            "priority_support",
        ],
    },
}


# ── License Models ───────────────────────────────────────────────────────────


class LicenseInfo(BaseModel):
    """License validation result returned to clients."""
    license_key: str = ""
    tier: LicenseTier = LicenseTier.COMMUNITY
    customer_name: str = ""
    customer_email: str = ""
    max_nodes: int = 0  # 0 = unlimited (community has unlimited nodes)
    features: List[str] = Field(default_factory=lambda: [
        "dashboard", "api_access", "job_routing",
        "amd_gpu_optimization", "one_time_probe", "credential_vault",
    ])
    issued_at: Optional[datetime] = None
    expires_at: Optional[datetime] = None
    valid: bool = True


class LicenseCreate(BaseModel):
    """Request body for creating a new license."""
    customer_name: str
    customer_email: str
    tier: LicenseTier = LicenseTier.PROFESSIONAL
    max_nodes: Optional[int] = None  # None = use tier default
    features: Optional[List[str]] = None  # None = use tier default
    duration_months: int = 12


class LicenseRecord(BaseModel):
    """Full license record as stored in the database."""
    license_id: int = 0
    customer_name: str = ""
    customer_email: str = ""
    license_key: str = ""
    tier: LicenseTier = LicenseTier.COMMUNITY
    max_nodes: int = 5
    features: List[str] = Field(default_factory=list)
    issued_at: datetime = Field(default_factory=datetime.utcnow)
    expires_at: Optional[datetime] = None
    active: bool = True


# ── Telemetry Models ─────────────────────────────────────────────────────────


class TelemetryPayload(BaseModel):
    """Heartbeat telemetry sent by the client."""
    node_count: int = 0
    hardware_summary: Dict[str, Any] = Field(default_factory=dict)
    mesh_version: str = ""
    uptime_hours: float = 0.0
    avg_cpu_pct: float = 0.0
    avg_memory_pct: float = 0.0
    total_jobs_run: int = 0
    total_retrains: int = 0


class UsageReport(BaseModel):
    """Usage statistics for a license."""
    license_key: str = ""
    daily_active_nodes: int = 0
    peak_nodes: int = 0
    total_heartbeats: int = 0
    feature_usage: Dict[str, int] = Field(default_factory=dict)
    last_heartbeat: Optional[datetime] = None
    telemetry_history: List[Dict[str, Any]] = Field(default_factory=list)
