"""
Given a printer_id and timestamp t, fetch health + operating conditions from
Supabase and return the PPO-recommended replacement schedule.
"""

import os
import numpy as np
from datetime import datetime, timezone
from supabase import create_client
from stable_baselines3 import PPO
from model import DegradationModel
from scheduling_rl import _ACTION_TABLE, COMPONENT_NAMES
from process_inputs import process_inputs

# ---------------------------------------------------------------------------
# Supabase client
# ---------------------------------------------------------------------------

_sb = create_client(os.environ["SUPABASE_URL"], os.environ["SUPABASE_KEY"])

# Column order must match INPUT_NAMES in model.py (C=9, obs vector is R^20)
_CONDITION_COLS = [
    "ambient_temperature_c",
    "build_chamber_temp_c",
    "ambient_humidity_pct",
    "powder_contamination_level",
    "print_hours",
    "build_volume_cm3",
    "recoating_speed_mm_s",
    "recoating_cycles",
    "maintenance_level",
]

_HEALTH_COLS = [
    "recoater_blade",
    "nozzle_plate",
    "heating_elements",
    "temperature_sensors",
    "insulation_panels",
    "firing_resistors",
    "cleaning_interface",
    "recoater_motor",
    "linear_rail",
]

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _as_aware(dt: datetime) -> datetime:
    """Return a timezone-aware datetime, treating naive datetimes as UTC."""
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt


def _parse_ts(s: str) -> datetime:
    """Parse an ISO timestamp string that may use a trailing Z."""
    return _as_aware(datetime.fromisoformat(s.replace("Z", "+00:00")))


# ---------------------------------------------------------------------------
# Fetch helpers
# ---------------------------------------------------------------------------

def _fetch_health(printer_id: str) -> np.ndarray:
    """Latest health snapshot for the printer."""
    row = (
        _sb.table("snapshots")
        .select(", ".join(_HEALTH_COLS))
        .eq("id", printer_id)
        .order("time_step_id", desc=True)
        .limit(1)
        .execute()
        .data
    )
    if not row:
        raise ValueError(f"No snapshot found for printer {printer_id}")
    return np.array([row[0][c] for c in _HEALTH_COLS], dtype=np.float64)


def _fetch_conditions(printer_id: str, t: datetime) -> np.ndarray:
    """Closest conditions row at or before t."""
    row = (
        _sb.table("conditions")
        .select(", ".join(_CONDITION_COLS))
        .eq("id", printer_id)
        .lte("timestamp", _as_aware(t).isoformat())
        .order("timestamp", desc=True)
        .limit(1)
        .execute()
        .data
    )
    if not row:
        raise ValueError(f"No conditions found for printer {printer_id} at {t}")
    # Coerce NULL columns to 0.0 (seed_data.py may omit some fields)
    return np.array([float(row[0][c] or 0.0) for c in _CONDITION_COLS], dtype=np.float64)

# ---------------------------------------------------------------------------
# Main prediction
# ---------------------------------------------------------------------------

def predict_replacements(
    printer_id: str,
    t: datetime,
    *,
    budget_remaining: float,
    W: float = 10_000.0,
    t_hours: float = 0.0,
    ppo_path: str = "scheduler_ppo",
    model_path: str = "model.npz",
) -> dict:
    DegradationModel.load(model_path)   # validates model exists
    ppo = PPO.load(ppo_path)

    health = _fetch_health(printer_id)
    X_t    = process_inputs(_fetch_conditions(printer_id, t))

    obs = np.concatenate([health, X_t, [budget_remaining / W], [t_hours]]).astype(np.float32)
    action, _ = ppo.predict(obs, deterministic=True)
    bits = _ACTION_TABLE[int(action)]

    to_replace = [COMPONENT_NAMES[i] for i, b in enumerate(bits) if b]
    return {
        "printer_id":    printer_id,
        "timestamp":     _as_aware(t).isoformat(),
        "health":        dict(zip(COMPONENT_NAMES, health.tolist())),
        "conditions":    dict(zip(_CONDITION_COLS, X_t.tolist())),
        "replace":       to_replace,
        "action_id":     int(action),
    }


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import json, sys
    printer_id = sys.argv[1] if len(sys.argv) > 1 else "printer_001"
    t = _parse_ts(sys.argv[2]) if len(sys.argv) > 2 else datetime.now(tz=timezone.utc)
    budget = float(sys.argv[3]) if len(sys.argv) > 3 else 10_000.0

    result = predict_replacements(printer_id, t, budget_remaining=budget)
    print(json.dumps(result, indent=2))