"""Shared data models for the Radix Edge Agent."""

from __future__ import annotations

import json
from dataclasses import dataclass, field, asdict
from enum import Enum
from typing import Any


class JobStatus(str, Enum):
    QUEUED = "queued"
    SCHEDULED = "scheduled"
    RUNNING = "running"
    SUCCEEDED = "succeeded"
    FAILED = "failed"
    CANCELLED = "cancelled"


class GPUStatus(str, Enum):
    IDLE = "idle"
    ALLOCATED = "allocated"
    ERROR = "error"


@dataclass
class Job:
    job_id: str
    user_id: str
    image: str
    status: JobStatus = JobStatus.QUEUED
    name: str | None = None
    command: list[str] | None = None
    args: list[str] | None = None
    env: dict[str, str] | None = None
    job_type: str = "generic"
    priority: int = 5
    num_gpus: int = 1
    gpu_mem_required_mb: int = 0
    timeout_seconds: int = 3600
    scheduled_gpus: list[int] | None = None
    container_id: str | None = None
    exit_code: int | None = None
    error_message: str | None = None
    stdout_tail: str | None = None
    stderr_tail: str | None = None
    runtime_seconds: float | None = None
    created_at: str | None = None
    scheduled_at: str | None = None
    started_at: str | None = None
    completed_at: str | None = None
    metadata: dict[str, Any] | None = None
    origin: str = "local"
    upstream_job_id: str | None = None

    @classmethod
    def from_row(cls, row) -> Job:
        """Create a Job from a sqlite3.Row."""
        d = dict(row)
        for json_field in ("command", "args", "env", "scheduled_gpus", "metadata"):
            if d.get(json_field) and isinstance(d[json_field], str):
                d[json_field] = json.loads(d[json_field])
        d["status"] = JobStatus(d["status"])
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})

    def to_dict(self) -> dict[str, Any]:
        """Serialize for API responses."""
        d = asdict(self)
        d["status"] = self.status.value
        return d


@dataclass
class GPU:
    gpu_index: int
    name: str
    memory_total_mb: int
    uuid: str | None = None
    memory_used_mb: int = 0
    memory_free_mb: int = 0
    utilization_pct: float = 0.0
    temperature_c: float = 0.0
    power_draw_w: float = 0.0
    power_limit_w: float = 0.0
    pcie_gen: int = 0
    nvlink_peers: list[int] | None = None
    assigned_job_id: str | None = None
    status: GPUStatus = GPUStatus.IDLE
    last_seen_at: str | None = None

    @classmethod
    def from_row(cls, row) -> GPU:
        d = dict(row)
        if d.get("nvlink_peers") and isinstance(d["nvlink_peers"], str):
            d["nvlink_peers"] = json.loads(d["nvlink_peers"])
        d["status"] = GPUStatus(d["status"])
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})

    def to_dict(self) -> dict[str, Any]:
        d = asdict(self)
        d["status"] = self.status.value
        return d


@dataclass
class ScoringResult:
    chosen_gpu: str
    priority_score: float
    gpu_indices: list[int] = field(default_factory=list)
    avoid_co_locate_with: list[str] = field(default_factory=list)
    terms: dict[str, float] = field(default_factory=dict)


@dataclass
class FEPDecision:
    latent_state: str
    state_probs: dict[str, float]
    observations: dict[str, float]
    chosen_action: str
    free_energy: float
    decided_at: str | None = None


@dataclass
class BrainParams:
    tau: float = 0.8
    alpha: float = 0.1
    ema_beta: float = 0.15
    w_queue_depth: float = 0.5
    w_gpu_saturation: float = 0.3
    w_node_pressure: float = 0.2
    updated_at: str | None = None

    @classmethod
    def from_row(cls, row) -> BrainParams:
        d = dict(row)
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)
