from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import os
import pandas as pd
import uvicorn

app = FastAPI()

try:
    # This works when running as a script (Hugging Face)
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
except NameError:
    # This works when running in a Notebook (Colab)
    BASE_DIR = os.getcwd()

model_path = os.path.join(BASE_DIR, "superkart_sales_rf_model_v1.joblib")

# Load the model
try:
    model = joblib.load(model_path)
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")

class PredictionRequest(BaseModel):
    Product_Id: str
    Product_Weight: float
    Product_Sugar_Content: str
    Product_Allocated_Area: float
    Product_Type: str
    Product_MRP: float
    Store_Establishment_Year: int 
    Store_Age: int
    Store_Size: str
    Store_Location_City_Type: str
    Store_Type: str

@app.get("/")
def health():
    return {"status": "ok", "message": "SuperKart API is Live"}

@app.post("/predict")
def predict(data: PredictionRequest):
    # 1. Convert Pydantic object to dictionary
    input_dict = data.model_dump()
    
    # 2. REPLICATE NOTEBOOK FEATURE ENGINEERING
    # Extract Product_Category_Type from Product_Id (e.g., 'FD6114' -> 'FD')
    input_dict['Product_Category_Type'] = input_dict['Product_Id'][:2]
    
    # Standardize Sugar Content
    sugar = input_dict['Product_Sugar_Content'].lower().replace('reg', 'regular').strip().title()
    input_dict['Product_Sugar_Content'] = sugar

    # 3. Create DataFrame
    df = pd.DataFrame([input_dict])
    
    # 4. Drop Product_Id as done in notebook
    df = df.drop('Product_Id', axis=1)
    
    # 5. Predict using the loaded pipeline
    pred = model.predict(df)[0]
    
    return {"prediction": float(pred)}

# CRITICAL: Hugging Face Spaces must run on port 7860
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)