from fastapi import FastAPI, UploadFile, File, HTTPException
import pandas as pd
from backend.core.orchestrator import Orchestrator

app = FastAPI()
orchestrator = Orchestrator()

@app.post("/analyze")
async def analyze_dataset(file: UploadFile = File(...), target_column: str = "target"):
    try:
        df = pd.read_csv(file.file)
        result = orchestrator.run(df, target_column)

        # Format response for frontend
        dataset_info = result.get("dataset_info", {})
        strategy = result.get("strategy", {})

        response = {
            "columns": list(df.columns),
            "dataTypes": dataset_info.get("data_types", {}),
            "risks": dataset_info.get("risks", []),
            "problemType": result.get("problem_type"),
            "confidence": strategy.get("confidence", 0),
            "strategy": strategy
        }
        return response
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/train")
async def train_model(file: UploadFile = File(...), target_column: str = "target"):
    try:
        df = pd.read_csv(file.file)
        result = orchestrator.run(df, target_column, train=True)

        # Ensure strategy is included in the response
        strategy = result.get("strategy", {})
        response = {
            "strategy": strategy,
            "metrics": result.get("metrics", {}),
            "model_path": result.get("model_path", "/path/to/model.pkl"),
            "training_time": result.get("training_time", 0),
            "model_id": result.get("model_id", "trained_model_123")
        }
        return response
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/explain")
async def explain_model(file: UploadFile = File(...), target_column: str = "target"):
    try:
        df = pd.read_csv(file.file)
        result = orchestrator.run(df, target_column, train=True)
        return {
            "strategy_explanation": result.get("strategy_explanation"),
            "metrics": result.get("metrics", {}),
            "feature_importance": result.get("feature_importance", [])
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/predict")
async def predict(data: dict):
    try:
        # Load the trained model
        model = orchestrator.model_io.load("exports/models/trained_model.pkl")
        # Prepare data for prediction
        df = pd.DataFrame([data])
        preds = model.predict(df)
        return {"prediction": preds.tolist()}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/health")
def health():
    return {"status": "ok"}