import os
from datetime import datetime
from flask import Flask, request, jsonify
from flask_cors import CORS
from dotenv import load_dotenv
import numpy as np

from predict_from_supabase import predict_replacements
from model import DegradationModel, COMPONENT_NAMES
from process_inputs import process_inputs

load_dotenv()

app = Flask(__name__)
CORS(app)

# Load the default degradation model once at startup
_DEFAULT_MODEL_PATH = "model.npz"
_degradation_model: DegradationModel | None = None

def _get_model(path: str = _DEFAULT_MODEL_PATH) -> DegradationModel:
    global _degradation_model
    if _degradation_model is None:
        _degradation_model = DegradationModel.load(path)
    return _degradation_model


@app.route("/api/predict", methods=["POST"])
def predict():
    data = request.json or {}

    h0_raw = data.get("h0")
    X_raw  = data.get("X")
    weeks  = float(data.get("weeks", 6))
    model_path = data.get("model_path", _DEFAULT_MODEL_PATH)

    if h0_raw is None:
        return jsonify({"error": "h0 (current health state) is required"}), 400
    if X_raw is None:
        return jsonify({"error": "X (input conditions) is required"}), 400

    try:
        model = _get_model(model_path)
    except FileNotFoundError:
        return jsonify({"error": f"model file not found: {model_path}"}), 404

    h0 = np.array(h0_raw, dtype=float)
    X  = process_inputs(np.array(X_raw, dtype=float))

    if h0.shape != (model.N,):
        return jsonify({"error": f"h0 must have {model.N} elements, got {h0.shape}"}), 400
    if X.shape != (model.C,):
        return jsonify({"error": f"X must have {model.C} elements, got {X.shape}"}), 400

    # Euler integration from h0, deterministic, 100 steps per week
    tau_hours = weeks * 7 * 24.0
    n_steps   = int(weeks) * 20
    dt        = tau_hours / n_steps
    steps_per_week = n_steps // int(weeks)

    h = h0.copy()
    weekly_snapshots = [h0.tolist()]

    for step in range(1, n_steps + 1):
        h = np.clip(h + dt * model.f(h, X), 0.0, 1.0)
        if step % steps_per_week == 0:
            weekly_snapshots.append(h.tolist())

    return jsonify({
        "component_names": COMPONENT_NAMES,
        "weeks": weeks,
        "weekly_health": weekly_snapshots,   # index 0 = now, index k = after k weeks
        "final_health": h.tolist(),
    })


@app.route("/api/schedule", methods=["POST"])
def schedule():
    data = request.json or {}

    printer_id       = data.get("printer_id", "").strip()
    timestamp_str    = data.get("timestamp", "").strip()
    budget_remaining = data.get("budget_remaining")
    t_hours          = data.get("t_hours", 0.0)

    if not printer_id:
        return jsonify({"error": "printer_id is required"}), 400
    if not timestamp_str:
        return jsonify({"error": "timestamp is required"}), 400
    if budget_remaining is None:
        return jsonify({"error": "budget_remaining is required"}), 400

    try:
        t = datetime.fromisoformat(timestamp_str)
    except ValueError:
        return jsonify({"error": f"invalid timestamp: {timestamp_str}"}), 400

    ppo_path   = f"scheduler_ppo"
    model_path = f"model.npz"

    if not os.path.exists(f"{ppo_path}.zip"):
        return jsonify({"error": f"no PPO model found for printer {printer_id}"}), 404
    if not os.path.exists(model_path):
        return jsonify({"error": f"no degradation model found for printer {printer_id}"}), 404

    try:
        result = predict_replacements(
            printer_id,
            t,
            budget_remaining=float(budget_remaining),
            t_hours=float(t_hours),
            ppo_path=ppo_path,
            model_path=model_path,
        )
    except ValueError as e:
        return jsonify({"error": str(e)}), 422

    return jsonify(result)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7861, debug=False)