import os
import uuid
import pandas as pd
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel, Field
from typing import Dict, Any, List, Optional
import io
import json
import sys

# Ensure backend directory is in path for imports when running from root
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from predictor import get_predictor

app = FastAPI(
    title="Suspicious Transaction & Mule Account Detector API",
    description="Backend API running ML predictions to identify suspicious transactions and mule accounts.",
    version="1.0.0"
)

# Allow CORS for React frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], # In development, allow all
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

def clean_nans(data):
    if isinstance(data, dict):
        return {k: clean_nans(v) for k, v in data.items()}
    elif isinstance(data, list):
        return [clean_nans(v) for v in data]
    elif isinstance(data, float):
        if np.isnan(data) or np.isinf(data):
            return None
        return data
    elif pd.isna(data): # handles pd.NaT, None, np.nan
        return None
    else:
        return data

# Directory to save processed runs
RUNS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "runs"))
os.makedirs(RUNS_DIR, exist_ok=True)

# Pydantic model for single prediction
# We accept a dynamic dict of fields to handle the thousands of columns.
class SingleTransactionRequest(BaseModel):
    data: Dict[str, Any]
    model_name: Optional[str] = "voting"
    threshold: Optional[float] = 0.30

class SingleTransactionResponse(BaseModel):
    probability: float
    is_suspicious: bool
    risk_level: str

@app.get("/api/health")
def health_check():
    try:
        get_predictor()
        return {"status": "healthy", "model_loaded": True}
    except Exception as e:
        return {"status": "degraded", "error": str(e), "model_loaded": False}

@app.get("/api/model-info")
def model_info():
    predictor = get_predictor()
    return {
        "cat_cols": predictor.cat_cols,
        "features_count": len(predictor.ohe_columns),
        "available_models": list(predictor.models.keys()),
        "default_thresholds": {
            "voting": 0.30,
            "xgboost": 0.30,
            "random_forest": 0.28
        }
    }

@app.post("/api/predict-single", response_model=SingleTransactionResponse)
def predict_single(req: SingleTransactionRequest):
    try:
        predictor = get_predictor()
        res = predictor.predict_transaction(
            data_row=req.data,
            model_name=req.model_name,
            threshold=req.threshold
        )
        return res
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

@app.post("/api/predict-csv")
async def predict_csv(
    file: UploadFile = File(...),
    model_name: str = Form("voting"),
    threshold: float = Form(0.30)
):
    # Validate file extension
    if not file.filename.endswith('.csv'):
        raise HTTPException(status_code=400, detail="Only CSV files are supported.")
    
    try:
        # Read uploaded file
        contents = await file.read()
        df = pd.read_csv(io.BytesIO(contents))
        
        if len(df) == 0:
            raise HTTPException(status_code=400, detail="The uploaded CSV file is empty.")
            
        predictor = get_predictor()
        
        # Run prediction
        results = predictor.predict_dataframe(df, model_name=model_name, threshold=threshold)
        
        # Join predictions back to the original df
        df_annotated = df.copy()
        df_annotated['fraud_probability'] = results['probability']
        df_annotated['is_suspicious'] = results['is_suspicious']
        df_annotated['risk_level'] = results['risk_level']
        
        # Generate session ID and save the annotated file
        session_id = str(uuid.uuid4())
        run_file_path = os.path.join(RUNS_DIR, f"{session_id}.csv")
        df_annotated.to_csv(run_file_path, index=False)
        
        # Calculate statistics
        total_count = len(df_annotated)
        suspicious_count = int(df_annotated['is_suspicious'].sum())
        suspicious_rate = float(df_annotated['is_suspicious'].mean())
        
        # Risk level counts
        risk_counts = df_annotated['risk_level'].value_counts().to_dict()
        risk_distribution = {
            "High": int(risk_counts.get("High", 0)),
            "Medium": int(risk_counts.get("Medium", 0)),
            "Low": int(risk_counts.get("Low", 0))
        }
        
        # Aggregate stats by categorical features for insights
        cat_analyses = {}
        for col in ['F3886', 'F3893', 'F3891', 'F3892']:
            if col in df_annotated.columns:
                # Group by column, get total and suspicious counts
                grouped = df_annotated.groupby(col).agg(
                    total=('fraud_probability', 'count'),
                    suspicious=('is_suspicious', 'sum'),
                    avg_prob=('fraud_probability', 'mean')
                ).reset_index()
                # Sort by suspicious desc
                grouped = grouped.sort_values(by='suspicious', ascending=False)
                # Convert to dict
                cat_analyses[col] = grouped.to_dict(orient='records')
                
        # Time-series analysis: group by date if F3888 is present
        time_series = []
        if 'F3888' in df_annotated.columns:
            # Parse dates to string YYYY-MM-DD
            dates = pd.to_datetime(df_annotated['F3888'], errors='coerce')
            df_temp = df_annotated.copy()
            df_temp['date_str'] = dates.dt.strftime('%Y-%m-%d')
            # drop NaT rows
            df_temp = df_temp.dropna(subset=['date_str'])
            if len(df_temp) > 0:
                ts_grouped = df_temp.groupby('date_str').agg(
                    total=('fraud_probability', 'count'),
                    suspicious=('is_suspicious', 'sum')
                ).reset_index()
                # Sort by date
                ts_grouped = ts_grouped.sort_values(by='date_str')
                time_series = ts_grouped.to_dict(orient='records')

        # Get preview of rows (first 100 rows)
        # We only return key display columns + prediction columns to keep payload small
        display_cols = ['Unnamed: 0', 'F3888', 'F3886', 'F3893', 'F3891', 'F3892', 'F3887', 'F3894', 'F3895']
        existing_display_cols = [c for c in display_cols if c in df_annotated.columns]
        extra_cols = ['fraud_probability', 'is_suspicious', 'risk_level']
        
        preview_df = df_annotated[existing_display_cols + extra_cols].head(100)
        # Fill NaN with None for JSON serialization
        preview_df = preview_df.replace({np.nan: None})
        preview_records = preview_df.to_dict(orient='records')
        
        # Get top suspicious rows (top 50 highest probability)
        top_suspicious_df = df_annotated[existing_display_cols + extra_cols].sort_values(by='fraud_probability', ascending=False).head(50)
        top_suspicious_df = top_suspicious_df.replace({np.nan: None})
        top_suspicious_records = top_suspicious_df.to_dict(orient='records')

        # Identify potential mule accounts based on recurring account IDs (e.g. F3887 or F3895)
        # A mule account candidate is an account with high risk average or multiple suspicious transactions
        mule_accounts = []
        acct_col = 'F3887' if 'F3887' in df_annotated.columns else ('F3895' if 'F3895' in df_annotated.columns else None)
        if acct_col:
            mule_grouped = df_annotated.groupby(acct_col).agg(
                tx_count=('fraud_probability', 'count'),
                suspicious_count=('is_suspicious', 'sum'),
                max_prob=('fraud_probability', 'max'),
                avg_prob=('fraud_probability', 'mean')
            ).reset_index()
            # Filter for accounts with at least 1 suspicious transaction or avg_prob > threshold
            mule_candidates = mule_grouped[mule_grouped['suspicious_count'] > 0].sort_values(by=['suspicious_count', 'avg_prob'], ascending=False).head(20)
            mule_candidates = mule_candidates.replace({np.nan: None})
            mule_accounts = mule_candidates.to_dict(orient='records')

        return clean_nans({
            "session_id": session_id,
            "total_count": total_count,
            "suspicious_count": suspicious_count,
            "suspicious_rate": round(suspicious_rate, 4),
            "risk_distribution": risk_distribution,
            "categorical_analyses": cat_analyses,
            "time_series": time_series,
            "preview": preview_records,
            "top_suspicious": top_suspicious_records,
            "mule_accounts": mule_accounts
        })
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"CSV processing failed: {str(e)}")

@app.get("/api/download-run/{session_id}")
def download_run(session_id: str):
    file_path = os.path.join(RUNS_DIR, f"{session_id}.csv")
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="Processed file not found. Session may have expired.")
    return FileResponse(
        path=file_path,
        media_type="text/csv",
        filename=f"suspicious_detection_report_{session_id[:8]}.csv"
    )

# Serve React static frontend in production if built
from fastapi.staticfiles import StaticFiles
static_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "static"))
if os.path.exists(static_dir):
    app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")