"""Pydantic models for mesh REST payloads."""
from __future__ import annotations

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

from pydantic import BaseModel, Field, field_validator


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

class NodeStatus(str, Enum):
    ONLINE = "online"
    DEGRADED = "degraded"
    OFFLINE = "offline"
    ELECTION = "election"


class JobStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class JobType(str, Enum):
    PROBE = "probe"
    BENCHMARK = "benchmark"
    TRAINING = "training"
    INFERENCE = "inference"
    CUSTOM = "custom"


# ── Hardware Models ──────────────────────────────────────────────────────────

class GPUInfo(BaseModel):
    name: str = ""
    vendor: str = ""  # "amd", "nvidia"
    vram_mb: int = 0
    arch: str = ""
    driver: str = ""
    temperature_c: float = 0.0
    utilization_pct: float = 0.0
    memory_used_mb: int = 0


class CPUInfo(BaseModel):
    model: str = ""
    cores: int = 0
    physical_cores: int = 0
    threads_per_core: int = 1
    freq_mhz: float = 0.0
    l1d_kb: int = 0
    l2_kb: int = 0
    l3_kb: int = 0
    numa_nodes: int = 1


class FPGAInfo(BaseModel):
    name: str = ""
    vendor: str = ""
    pcie_bdf: str = ""
    driver: str = ""


class MountInfo(BaseModel):
    mountpoint: str = ""
    device: str = ""
    fstype: str = ""
    total_mb: int = 0
    free_mb: int = 0
    seq_read_mbps: float = 0.0
    seq_write_mbps: float = 0.0


class HardwareInventory(BaseModel):
    hostname: str = ""
    platform: str = ""
    cpu: CPUInfo = Field(default_factory=CPUInfo)
    memory_total_mb: int = 0
    gpus: List[GPUInfo] = Field(default_factory=list)
    fpgas: List[FPGAInfo] = Field(default_factory=list)
    has_pytorch: bool = False
    has_rocm: bool = False
    has_cuda: bool = False
    # Slento CUDA (CUDA→ROCm translation layer)
    has_slento_cuda: bool = False
    slento_cuda_version: str = ""
    slento_cuda_gpu_archs: List[str] = Field(default_factory=list)
    slento_cuda_lib_path: str = ""
    # Storage mounts
    mounts: List[MountInfo] = Field(default_factory=list)
    scratch_paths: List[str] = Field(default_factory=list)


# ── Health Models ────────────────────────────────────────────────────────────

class HealthSnapshot(BaseModel):
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    cpu_pct: float = 0.0
    cpu_freq_mhz: float = 0.0
    memory_used_pct: float = 0.0
    memory_available_mb: int = 0
    disk_used_pct: float = 0.0
    load_1m: float = 0.0
    load_5m: float = 0.0
    load_15m: float = 0.0
    gpu_utilizations: List[float] = Field(default_factory=list)
    gpu_temperatures: List[float] = Field(default_factory=list)
    gpu_memory_used_mb: List[int] = Field(default_factory=list)
    net_bytes_sent: int = 0
    net_bytes_recv: int = 0
    uptime_s: float = 0.0
    # Storage: per-mount free space (mountpoint -> free_mb)
    mount_free_mb: Dict[str, int] = Field(default_factory=dict)
    # Network: measured bandwidth and per-peer latency
    net_bandwidth_mbps: float = 0.0
    lan_latency_ms: Dict[str, float] = Field(default_factory=dict)
    # Dataset locality: IDs of datasets cached on this node
    cached_datasets: List[str] = Field(default_factory=list)


# ── Node Models ──────────────────────────────────────────────────────────────

class NodeRegistration(BaseModel):
    node_id: str
    hostname: str
    agent_url: str
    hardware: HardwareInventory
    tags: List[str] = Field(default_factory=list)
    # WAN: True if agent is in NAT mode (push-only, no inbound connections).
    # Controller will not attempt to dispatch jobs or trigger probes on this node.
    nat_mode: bool = False

    @field_validator("hostname")
    @classmethod
    def validate_hostname(cls, v: str) -> str:
        from mesh.security.validation import sanitize_hostname
        return sanitize_hostname(v)

    @field_validator("agent_url")
    @classmethod
    def validate_agent_url(cls, v: str) -> str:
        from urllib.parse import urlparse
        parsed = urlparse(v.strip())
        if parsed.scheme not in ("http", "https"):
            raise ValueError("agent_url must use http or https scheme")
        if not parsed.hostname:
            raise ValueError("agent_url must have a hostname")
        if parsed.port is not None and (parsed.port < 1 or parsed.port > 65535):
            raise ValueError(f"agent_url port out of range: {parsed.port}")
        return v.strip()


class NodeHeartbeat(BaseModel):
    node_id: str
    health: HealthSnapshot
    active_jobs: int = 0


class NodeInfo(BaseModel):
    node_id: str
    hostname: str
    agent_url: str
    status: NodeStatus = NodeStatus.ONLINE
    hardware: HardwareInventory = Field(default_factory=HardwareInventory)
    last_health: Optional[HealthSnapshot] = None
    last_heartbeat: Optional[datetime] = None
    registered_at: Optional[datetime] = None
    tags: List[str] = Field(default_factory=list)
    # WAN: True if this node is behind NAT (push-only). The controller cannot
    # initiate connections to it — jobs must be picked up via polling.
    nat_mode: bool = False


# ── Job Models ───────────────────────────────────────────────────────────────

class JobSubmission(BaseModel):
    job_id: Optional[str] = None  # Set by controller for dispatched jobs
    job_type: JobType = JobType.CUSTOM
    command: str = Field(default="", max_length=10000)
    args: Dict[str, Any] = Field(default_factory=dict)
    target_node: Optional[str] = None  # None = auto-route
    priority: int = 0
    timeout_s: float = 3600.0
    # Data-aware routing hints
    input_size_mb: float = 0.0
    output_size_mb: float = 0.0
    dataset_ids: List[str] = Field(default_factory=list)
    # Runtime requirements
    requires_cuda: bool = False
    requires_rocm: bool = False
    requires_slento_cuda: bool = False  # True = can run on AMD via Slento CUDA
    min_gpu_vram_mb: int = 0
    min_free_disk_mb: int = 0
    scratch_space_mb: int = 0
    # Network sensitivity
    max_acceptable_latency_ms: float = 0.0  # 0 = no constraint

    @field_validator("command")
    @classmethod
    def validate_command(cls, v: str) -> str:
        if "\x00" in v:
            raise ValueError("Command must not contain null bytes")
        return v


class JobInfo(BaseModel):
    job_id: str = ""
    job_type: JobType = JobType.CUSTOM
    command: str = ""
    args: Dict[str, Any] = Field(default_factory=dict)
    status: JobStatus = JobStatus.PENDING
    assigned_node: Optional[str] = None
    submitted_at: Optional[datetime] = None
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    priority: int = 0


class JobResult(BaseModel):
    job_id: str
    status: JobStatus
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    duration_s: float = 0.0


# ── Atlas / JEPA Models ─────────────────────────────────────────────────────

_ATLAS_DB_NAME_RE = re.compile(r"^[a-zA-Z0-9_\.]+$")


class AtlasSyncPayload(BaseModel):
    node_id: str
    atlas_db_name: str
    data_points: List[Dict[str, Any]] = Field(default_factory=list, max_length=10000)
    invariants: List[Dict[str, Any]] = Field(default_factory=list, max_length=10000)

    @field_validator("atlas_db_name")
    @classmethod
    def validate_atlas_db_name(cls, v: str) -> str:
        if not v or not _ATLAS_DB_NAME_RE.match(v):
            raise ValueError(
                "atlas_db_name must contain only alphanumeric characters, underscores, and dots"
            )
        if ".." in v:
            raise ValueError("atlas_db_name must not contain '..'")
        return v

    @field_validator("data_points")
    @classmethod
    def validate_data_points_count(cls, v: list) -> list:
        if len(v) > 10000:
            raise ValueError(f"data_points exceeds maximum of 10000 items ({len(v)} given)")
        return v

    @field_validator("invariants")
    @classmethod
    def validate_invariants_count(cls, v: list) -> list:
        if len(v) > 10000:
            raise ValueError(f"invariants exceeds maximum of 10000 items ({len(v)} given)")
        return v


class JEPAStats(BaseModel):
    trained: bool = False
    online_updates: int = 0
    memory_bank_size: int = 0
    training_samples: int = 0
    mean_error: float = 0.0
    confidence_by_regime: Dict[str, float] = Field(default_factory=dict)
    last_retrain: Optional[datetime] = None


class JEPAFeedback(BaseModel):
    workload: Dict[str, float]
    config: Dict[str, float]
    actual_performance: Dict[str, float]
    machine_name: Optional[str] = None


# ── System Optimization Models ──────────────────────────────────────────────

class SystemTuning(BaseModel):
    """System-level tuning recommendations generated from atlas/JEPA data."""
    cpu_governor: str = "performance"
    hugepages_2m: int = 0
    transparent_hugepages: str = "madvise"
    numa_balancing: bool = True
    swappiness: int = 10
    dirty_ratio: int = 20
    dirty_background_ratio: int = 5
    vfs_cache_pressure: int = 50
    sched_autogroup: bool = False
    sched_util_clamp_min: int = 0
    sched_util_clamp_max: int = 1024
    sched_rr_timeslice_ms: int = 100
    io_schedulers: Dict[str, str] = Field(default_factory=dict)  # device -> scheduler
    net_rmem_max: int = 16777216
    net_wmem_max: int = 16777216
    net_tcp_rmem: str = "4096 131072 16777216"
    net_tcp_wmem: str = "4096 65536 16777216"
    net_somaxconn: int = 65535
    irq_affinity: Dict[str, int] = Field(default_factory=dict)  # nic -> numa_node
    applied_at: Optional[datetime] = None
    node_id: Optional[str] = None


class StorageProbeResult(BaseModel):
    device: str = ""
    mountpoint: str = ""
    seq_read_mbps: float = 0.0
    seq_write_mbps: float = 0.0
    rand_read_iops: float = 0.0
    rand_write_iops: float = 0.0
    rand_read_latency_us: float = 0.0
    fsync_mean_us: float = 0.0
    fsync_p99_us: float = 0.0
    optimal_queue_depth: int = 1


class NetworkProbeResult(BaseModel):
    interface: str = ""
    speed_mbps: int = 0
    numa_node: int = -1
    driver: str = ""
    tcp_throughput_mbps: float = 0.0
    latency_ms: float = 0.0


# ── Response Models ──────────────────────────────────────────────────────────

class StatusResponse(BaseModel):
    status: str = "ok"
    message: str = ""
    data: Optional[Dict[str, Any]] = None
