import json
import joblib
import numpy as np
import pandas as pd

from pathlib import Path
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import Dict, Any


ARTIFACT_DIR = Path(__file__).resolve().parent

ml_assets = {}


class PredictRequest(BaseModel):
    features: Dict[str, Any] = Field(..., description="Feature dict untuk 1 well/design row")


def encode_for_inference(
    df_infer: pd.DataFrame,
    trained_feature_columns,
    encoded_category_maps,
) -> pd.DataFrame:
    X = df_infer.copy()

    # encode kategorikal sesuai mapping training
    for col, cmap in encoded_category_maps.items():
        if col in X.columns:
            X[col] = X[col].map(cmap)
            # unknown category -> -1
            X[col] = X[col].fillna(-1)

    # tambah kolom yang hilang
    for col in trained_feature_columns:
        if col not in X.columns:
            X[col] = 0

    # buang kolom ekstra, lalu reorder persis seperti training
    X = X[trained_feature_columns].copy()

    # paksa numerik
    for col in X.columns:
        X[col] = pd.to_numeric(X[col], errors="coerce").fillna(0)

    return X


def enrich_derived_features(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()

    # safe pressure limit
    required_pressure_cols = [
        "max_allowable_surface_pressure_psi",
        "casing_pressure_limit_psi",
        "tubing_pressure_limit_psi",
    ]
    if all(c in out.columns for c in required_pressure_cols):
        out["safe_pressure_limit_psi"] = out[required_pressure_cols].min(axis=1)

    if "safe_pressure_limit_psi" in out.columns and "max_pressure_psi" in out.columns:
        out["pressure_headroom_psi"] = out["safe_pressure_limit_psi"] - out["max_pressure_psi"]

    if "avg_planned_rate_bpm" in out.columns and "max_pump_rate_bpm" in out.columns:
        denom = out["max_pump_rate_bpm"].replace(0, np.nan)
        out["rate_to_pump_capacity"] = (out["avg_planned_rate_bpm"] / denom).fillna(0)

    if "inventory_proppant_ton" in out.columns and "total_planned_proppant_ton" in out.columns:
        denom = out["total_planned_proppant_ton"].replace(0, np.nan)
        out["inventory_proppant_coverage"] = (out["inventory_proppant_ton"] / denom).fillna(0)

    if "inventory_fluid_bbl" in out.columns and "total_planned_fluid_bbl" in out.columns:
        denom = out["total_planned_fluid_bbl"].replace(0, np.nan)
        out["inventory_fluid_coverage"] = (out["inventory_fluid_bbl"] / denom).fillna(0)

    if "total_planned_fluid_bbl" in out.columns and "total_planned_proppant_ton" in out.columns:
        denom = out["total_planned_proppant_ton"].replace(0, np.nan)
        out["fluid_to_proppant_ratio"] = (out["total_planned_fluid_bbl"] / denom).fillna(0)

    if "planned_stage_count" in out.columns and "total_planned_proppant_ton" in out.columns:
        denom = out["planned_stage_count"].replace(0, np.nan)
        out["proppant_per_stage_ton"] = (out["total_planned_proppant_ton"] / denom).fillna(0)

    if "planned_stage_count" in out.columns and "total_planned_fluid_bbl" in out.columns:
        denom = out["planned_stage_count"].replace(0, np.nan)
        out["fluid_per_stage_bbl"] = (out["total_planned_fluid_bbl"] / denom).fillna(0)

    if "lateral_treated_length_m" in out.columns and "planned_stage_count" in out.columns:
        denom = out["planned_stage_count"].replace(0, np.nan)
        out["derived_stage_length_m"] = (out["lateral_treated_length_m"] / denom).fillna(0)

    return out


@asynccontextmanager
async def lifespan(app: FastAPI):
    with open(ARTIFACT_DIR / "artifact_manifest.json", "r", encoding="utf-8") as f:
        manifest = json.load(f)

    ml_assets["manifest"] = manifest
    ml_assets["uplift_model"] = joblib.load(ARTIFACT_DIR / "uplift_model.joblib")
    ml_assets["cost_model"] = joblib.load(ARTIFACT_DIR / "cost_model.joblib")
    ml_assets["risk_model"] = joblib.load(ARTIFACT_DIR / "risk_model.joblib")
    yield
    ml_assets.clear()


app = FastAPI(
    title="TRIDENT Design Recommender Inference API",
    version="1.0.0",
    lifespan=lifespan,
)


@app.get("/health")
def health():
    return {
        "status": "ok",
        "artifacts_loaded": bool(ml_assets),
        "model_version": ml_assets.get("manifest", {}).get("artifact_version"),
    }


@app.post("/predict")
def predict(req: PredictRequest):
    try:
        manifest = ml_assets["manifest"]
        trained_feature_columns = manifest["trained_feature_columns"]
        encoded_category_maps = manifest["encoded_category_maps"]
        default_eval_days = manifest["default_eval_days"]
        risk_penalty_usd = manifest["risk_penalty_usd"]

        df = pd.DataFrame([req.features])
        df = enrich_derived_features(df)

        X = encode_for_inference(
            df_infer=df,
            trained_feature_columns=trained_feature_columns,
            encoded_category_maps=encoded_category_maps,
        )

        uplift = float(ml_assets["uplift_model"].predict(X)[0])
        cost = float(ml_assets["cost_model"].predict(X)[0])
        risk = float(ml_assets["risk_model"].predict_proba(X)[0, 1])

        oil_price = float(df["realized_oil_price_usd_bbl"].iloc[0]) if "realized_oil_price_usd_bbl" in df.columns else 65.0
        gross_value = uplift * default_eval_days * oil_price
        design_score = gross_value - cost - risk * risk_penalty_usd

        response = {
            "pred_uplift_bopd": uplift,
            "pred_cost_usd": cost,
            "pred_screenout_risk": risk,
            "pred_gross_value_usd": gross_value,
            "pred_design_score_usd": design_score,
        }

        # optional info untuk debugging
        if "safe_pressure_limit_psi" in df.columns:
            response["safe_pressure_limit_psi"] = float(df["safe_pressure_limit_psi"].iloc[0])
        if "pressure_headroom_psi" in df.columns:
            response["pressure_headroom_psi"] = float(df["pressure_headroom_psi"].iloc[0])

        return response

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))