app/routers/predict.py

import io
import json
import math
import base64
import pickle
import pandas as pd
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form
from sqlalchemy.orm import Session

from app.database import get_db
from app.dependencies import (
    get_current_user, check_rate_limit, check_prediction_quota,
    track_usage, get_usage_summary,
)
from app.models.user import User
from app.models.mlmodel import MLModel
from app.schemas import PredictionInput, PredictionResponse, SinglePrediction, SmartPredictInput
from app.services.field_mapping import normalize_rows
from app.services.scoring import _churn_factor, _lead_factor
from app.services.training import predict_with_model
from app.services.benchmarking import update_benchmarks, compare_to_benchmark

router = APIRouter(prefix="/v1/predict", tags=["Prediction"])


def _apply_heuristics(data: list, model_type: str) -> list:
    """Apply heuristic rules to a list of data dicts."""
    results = []
    rule_func = _churn_factor if model_type == "churn" else _lead_factor

    for i, row in enumerate(data):
        score, factors = rule_func(row)

        if score >= 70:
            risk = "High Risk" if model_type == "churn" else "Hot"
            action = ("Immediate outreach + retention offer" if model_type == "churn"
                      else "Call immediately — high intent signals")
        elif score >= 40:
            risk = "Medium Risk" if model_type == "churn" else "Warm"
            action = ("Monitor + engagement campaign" if model_type == "churn"
                      else "Send case study + schedule demo")
        else:
            risk = "Low Risk" if model_type == "churn" else "Cold"
            action = ("No action needed" if model_type == "churn"
                      else "Add to email drip sequence")

        cid = row.get("customer_id") or row.get("lead_id") or row.get("id") or None

        results.append(SinglePrediction(
            index=i,
            score=float(score),
            risk_level=risk,
            recommended_action=action,
            risk_factors=factors,
            scoring_mode="heuristic",
            customer_id=str(cid) if cid else None,
            input_fields=row,
        ))

    return results


def _score_with_custom_model(ml_model, data: list, model_type: str) -> list:
    """Score rows with a trained model, mirroring the heuristic output shape."""
    scores = predict_with_model(
        ml_model.model_binary,
        ml_model.get_feature_names(),
        ml_model.encoders_json,
        data,
    )
    results = []
    for i, (row, score) in enumerate(zip(data, scores)):
        cid = row.get("customer_id") or row.get("lead_id") or row.get("id")
        if score >= 70:
            risk = "High Risk" if model_type == "churn" else "Hot"
            action = ("Immediate outreach + retention offer" if model_type == "churn"
                      else "Call immediately — high intent signals")
        elif score >= 40:
            risk = "Medium Risk" if model_type == "churn" else "Warm"
            action = ("Monitor + engagement campaign" if model_type == "churn"
                      else "Send case study + schedule demo")
        else:
            risk = "Low Risk" if model_type == "churn" else "Cold"
            action = ("No action needed" if model_type == "churn"
                      else "Add to email drip sequence")
        results.append(SinglePrediction(
            index=i, score=score, risk_level=risk,
            recommended_action=action, risk_factors=[],
            scoring_mode="custom_ml",
            customer_id=str(cid) if cid else None,
            input_fields=row,
        ))
    return results


@router.post("/smart", response_model=PredictionResponse)
def predict_smart(
    body: SmartPredictInput,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """Universal connector — send rows with ANY column names. We auto-map them
    to RevAI's signals (and derive durations from dates), then score. The
    response includes a `field_mapping` report of what was matched/missed."""
    if body.model_type not in ("churn", "lead"):
        raise HTTPException(status_code=400, detail="model_type must be 'churn' or 'lead'")

    rows, report = normalize_rows(body.data, body.mapping, body.model_type)

    n_predictions = len(rows)
    check_rate_limit(user, db, f"predict/{body.model_type}", n_predictions)
    check_prediction_quota(user, db, n_predictions)

    if body.model_id:
        ml_model = db.query(MLModel).filter(
            MLModel.id == body.model_id, MLModel.user_id == user.id
        ).first()
        if not ml_model:
            raise HTTPException(status_code=404, detail="Model not found")
        results = _score_with_custom_model(ml_model, rows, body.model_type)
        model_label = f"custom_ml_{body.model_id[:8]}"
    else:
        results = _apply_heuristics(rows, body.model_type)
        model_label = "heuristic"

    track_usage(user, db, f"predict/{body.model_type}", n_predictions)

    all_scores = [p.score for p in results]
    update_benchmarks(db, all_scores, body.model_type)
    benchmark = compare_to_benchmark(all_scores, body.model_type, db)

    return PredictionResponse(
        predictions=results,
        model_used=model_label,
        usage=get_usage_summary(user, db),
        benchmark=benchmark,
        field_mapping=report,
    )


@router.post("/csv", response_model=PredictionResponse)
async def predict_csv(
    file: UploadFile = File(..., description="CSV export of your customers/leads"),
    model_type: str = Form("churn", description="'churn' or 'lead'"),
    model_id: str = Form(None, description="Optional trained model ID"),
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """Upload a CSV instead of hand-building JSON. One column per signal,
    one row per customer/lead — we parse it and score every row."""
    if not file.filename or not file.filename.lower().endswith(".csv"):
        raise HTTPException(status_code=400, detail="Please upload a .csv file")
    if model_type not in ("churn", "lead"):
        raise HTTPException(status_code=400, detail="model_type must be 'churn' or 'lead'")

    raw = await file.read()
    try:
        df = pd.read_csv(io.BytesIO(raw))
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Could not parse CSV: {e}")

    df = df.dropna(how="all")
    if len(df) == 0:
        raise HTTPException(status_code=400, detail="CSV has no data rows")

    data = df.to_dict(orient="records")
    # pandas emits NaN for blanks; convert to None so scoring uses defaults
    for row in data:
        for k, v in list(row.items()):
            if isinstance(v, float) and math.isnan(v):
                row[k] = None

    n_predictions = len(data)
    check_rate_limit(user, db, f"predict/{model_type}", n_predictions)
    check_prediction_quota(user, db, n_predictions)

    if model_id:
        ml_model = db.query(MLModel).filter(
            MLModel.id == model_id, MLModel.user_id == user.id
        ).first()
        if not ml_model:
            raise HTTPException(status_code=404, detail="Model not found")
        results = _score_with_custom_model(ml_model, data, model_type)
        model_label = f"custom_ml_{model_id[:8]}"
    else:
        results = _apply_heuristics(data, model_type)
        model_label = "heuristic"

    track_usage(user, db, f"predict/{model_type}", n_predictions)

    all_scores = [p.score for p in results]
    update_benchmarks(db, all_scores, model_type)
    benchmark = compare_to_benchmark(all_scores, model_type, db)

    return PredictionResponse(
        predictions=results,
        model_used=model_label,
        usage=get_usage_summary(user, db),
        benchmark=benchmark,
    )


@router.post("/churn", response_model=PredictionResponse)
def predict_churn(
    body: PredictionInput,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    n_predictions = len(body.data)
    check_rate_limit(user, db, "predict/churn", n_predictions)
    used, limit = check_prediction_quota(user, db, n_predictions)

    if body.model_id:
        # Use custom trained model
        ml_model = db.query(MLModel).filter(
            MLModel.id == body.model_id, MLModel.user_id == user.id
        ).first()
        if not ml_model:
            raise HTTPException(status_code=404, detail="Model not found")

        scores = predict_with_model(
            ml_model.model_binary,
            ml_model.get_feature_names(),
            ml_model.encoders_json,
            body.data,
        )

        results = []
        for i, (row, score) in enumerate(zip(body.data, scores)):
            cid = row.get("customer_id") or row.get("id")
            if score >= 70:
                risk = "High Risk"
                action = "Immediate outreach + retention offer"
            elif score >= 40:
                risk = "Medium Risk"
                action = "Monitor + engagement campaign"
            else:
                risk = "Low Risk"
                action = "No action needed"

            results.append(SinglePrediction(
                index=i, score=score, risk_level=risk,
                recommended_action=action, risk_factors=[],
                scoring_mode="custom_ml",
                customer_id=str(cid) if cid else None,
                input_fields=row,
            ))

        model_label = f"custom_ml_{body.model_id[:8]}"
    else:
        results = _apply_heuristics(body.data, "churn")
        model_label = "heuristic"

    track_usage(user, db, "predict/churn", n_predictions)

    # ── Anonymous benchmark update + comparison ──
    all_scores = [p.score for p in results]
    update_benchmarks(db, all_scores, "churn")
    benchmark = compare_to_benchmark(all_scores, "churn", db)

    return PredictionResponse(
        predictions=results,
        model_used=model_label,
        usage=get_usage_summary(user, db),
        benchmark=benchmark,
    )


@router.post("/lead", response_model=PredictionResponse)
def predict_lead(
    body: PredictionInput,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    n_predictions = len(body.data)
    check_rate_limit(user, db, "predict/lead", n_predictions)
    used, limit = check_prediction_quota(user, db, n_predictions)

    if body.model_id:
        ml_model = db.query(MLModel).filter(
            MLModel.id == body.model_id, MLModel.user_id == user.id
        ).first()
        if not ml_model:
            raise HTTPException(status_code=404, detail="Model not found")

        scores = predict_with_model(
            ml_model.model_binary,
            ml_model.get_feature_names(),
            ml_model.encoders_json,
            body.data,
        )

        results = []
        for i, (row, score) in enumerate(zip(body.data, scores)):
            cid = row.get("lead_id") or row.get("id")
            if score >= 70:
                risk = "Hot"
                action = "Call immediately — high intent signals"
            elif score >= 40:
                risk = "Warm"
                action = "Send case study + schedule demo"
            else:
                risk = "Cold"
                action = "Add to email drip sequence"

            results.append(SinglePrediction(
                index=i, score=score, risk_level=risk,
                recommended_action=action, risk_factors=[],
                scoring_mode="custom_ml",
                customer_id=str(cid) if cid else None,
                input_fields=row,
            ))

        model_label = f"custom_ml_{body.model_id[:8]}"
    else:
        results = _apply_heuristics(body.data, "lead")
        model_label = "heuristic"

    track_usage(user, db, "predict/lead", n_predictions)

    all_scores = [p.score for p in results]
    update_benchmarks(db, all_scores, "lead")
    benchmark = compare_to_benchmark(all_scores, "lead", db)

    return PredictionResponse(
        predictions=results,
        model_used=model_label,
        usage=get_usage_summary(user, db),
        benchmark=benchmark,
    )