import os

import joblib
from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download
import pandas as pd
import numpy as np

# Local paths (used when these files exist, e.g. during local development)
model_uri = "../../../ml/models/FD004GBM_BAG"
data_uri = "../../../data/FE_data/test_FD004"

# Hugging Face Hub locations (used when the local files are not found)
HF_MODEL_REPO = "ZephyrNTempest/JetBoost-GBM"
HF_MODEL_FILE = "models/FD004GBM_BAG"
HF_DATA_REPO = "ZephyrNTempest/C-MAPSS"
HF_DATA_FILE = "data/FE_data/test_FD004"
HF_DATA_REPO_TYPE = (
    "dataset"  # "dataset" or "model", whichever you uploaded the data to
)


def resolve_path(local_path, repo_id, filename, repo_type="model"):
    """Return a local file path, downloading from the Hugging Face Hub if needed.

    Args:
        local_path (str): Local path to use if the file exists.
        repo_id (str): Hub repo id, e.g. "username/repo".
        filename (str): File name inside the Hub repo.
        repo_type (str): "model" or "dataset".

    Returns:
        str: Path to a readable local copy of the file.
    """
    if os.path.exists(local_path):
        return local_path
    return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type)


# Load the model
models = joblib.load(resolve_path(model_uri, HF_MODEL_REPO, HF_MODEL_FILE))
app = FastAPI()

# Allow the GitHub Pages site (and local dev) to call this API.
# Set ALLOWED_ORIGINS as a comma-separated list in the Space's variables.
allowed_origins = os.getenv(
    "ALLOWED_ORIGINS", "http://localhost:5173,https://YOUR_USERNAME.github.io"
).split(",")
app.add_middleware(
    CORSMiddleware,
    allow_origins=[o.strip() for o in allowed_origins],
    allow_methods=["GET"],
    allow_headers=["*"],
)


# Prune test set features
def prune_features(test, drop_pc2=False):
    """Drop min/last/max/delta feature columns (and pc2 if asked) from train and test."""
    patterns = ["min", "last", "max", "delta"]
    if drop_pc2:
        patterns.append("pc2")

    to_drop = [c for c in test.columns if any(p in c for p in patterns)]

    return test.drop(columns=to_drop, errors="ignore")


df = pd.read_parquet(
    resolve_path(data_uri, HF_DATA_REPO, HF_DATA_FILE, HF_DATA_REPO_TYPE)
)
df = df.drop(columns="RUL")
df = df.set_index("Engine")
df = prune_features(df)


# Average the preditions of the bagged models
class Bag:
    def __init__(self, models):
        self.models = models

    def predict(self, X):
        return np.mean([m.predict(X) for m in self.models], axis=0)


bag = Bag(models)


@app.get("/predict")
def inference(target_idxs: list[int] = Query(...), cycles: list[int] = Query(...)):
    """Return FD004GBM model predictions for a batch of engines,
    each at its own time cycle.

    Args:
        target_idxs (list[int]): Engine indices (e.g. 1..100) to predict on.
        cycles (list[int]): Time cycle for each engine. cycles[i] belongs to
            target_idxs[i], so both lists must be the same length.

    Returns:
        dict: {"Predicted RUL": [...]} with one prediction per engine,
            in the same order as target_idxs.
    """

    lookup = dict(zip(target_idxs, cycles))
    wanted = df.index.map(lookup)
    targets = df[df["Time"] == wanted]
    targets = targets.loc[target_idxs]
    rul_pred = bag.predict(targets)

    return {"Predicted RUL": rul_pred.tolist()}