"""
FastAPI Service for NASA SOH Prediction Model

Production API endpoint for battery health predictions.
Connects Python ML pipeline to React frontend apps.
"""

from fastapi import FastAPI, File, UploadFile, HTTPException, Depends, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from typing import List, Optional, Dict
from sqlalchemy.orm import Session
from datetime import datetime
import pandas as pd
import numpy as np
import pickle
import os
import base64
from pathlib import Path
import json
import warnings
warnings.filterwarnings('ignore')

from predict_soh import predict_soh, calculate_data_quality
from feature_engineering_causal import engineer_features_causal
from database import get_db, init_db, Asset, HealthPrediction

# Initialize FastAPI app
app = FastAPI(
    title="ZenVolt NASA SOH Trust Engine API",
    description="Production API for battery State of Health prediction using NASA-trained models",
    version="2.0.0"
)

# Initialize database on startup
@app.on_event("startup")
async def startup_event():
    init_db()
    print("✓ API service started")

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Request/Response Models
class BMSDataPoint(BaseModel):
    """Single BMS data point."""
    cycle_index: int
    voltage: float
    current: float
    temperature: float
    capacity: float
    timestamp: Optional[str] = None


class BMSDataInput(BaseModel):
    """BMS data input for prediction."""
    data: List[BMSDataPoint]
    initial_capacity: Optional[float] = None
    forecast_horizons: Optional[List[int]] = Field(default=[10, 50, 100])


class HealthPredictionResponse(BaseModel):
    """ML prediction response."""
    current_soh: Dict[str, float]
    health_score: float
    remaining_life: Dict[str, float]
    degradation_rate: Dict[str, float]
    forecasts: Dict[str, Dict[str, float]]
    data_quality: Dict[str, float]
    current_cycle: int
    confidence: float
    model_info: Dict[str, str]


class ModelInfoResponse(BaseModel):
    """Model information for Trust Engine proof."""
    model_type: str
    training_data: str
    validation_data: str
    metrics: Dict[str, float]
    feature_count: int
    version: str


class AssetCreate(BaseModel):
    """Model for creating a new asset."""
    owner_id: str
    project_name: str
    asset_id: str
    location: str
    capacity: float
    bms_provider: Optional[str] = None
    category: List[str] = Field(default_factory=list)
    project_stage: str = "Pending"
    initial_prediction: Optional[Dict] = None  # Initial prediction data


class AssetResponse(BaseModel):
    """Asset response model."""
    id: str
    owner_id: str
    projectName: str
    assetId: str
    location: str
    capacity: float
    bmsProvider: Optional[str]
    category: List[str]
    projectStage: str
    healthScore: Optional[float]
    remainingLife: Optional[float]
    degradationRate: Optional[float]
    confidence: Optional[float]
    stateOfCharge: Optional[float]
    voltage: Optional[float]
    temperature: Optional[float]
    isOnline: bool
    nftHash: Optional[str]
    nftMintDate: Optional[str]
    blockExplorerUrl: Optional[str]
    createdAt: Optional[str]
    updatedAt: Optional[str]
    lastUpdated: Optional[str]


class FleetMetricsResponse(BaseModel):
    """Fleet metrics response."""
    totalAssets: int
    avgHealth: float
    totalCapacity: float
    certifiedCount: int
    criticalCount: int
    avgRemainingLife: float


# Include asset routes safely
print("🔍 Attempting to mount asset routes...")
try:
    import asset_routes
    app.include_router(asset_routes.router)
    print("✅ SUCCESS: Asset routes mounted at /api/assets")
except Exception as e:
    print(f"❌ CRITICAL FAILURE: Could not mount asset routes: {e}")
    traceback.print_exc()


@app.get("/")
async def root():
    """Root endpoint."""
    environment = "Docker/HuggingFace" if os.path.exists("/.dockerenv") or os.environ.get("HOSTNAME", "").endswith("-space") else "Local/Standalone"
    return {
        "service": "ZenVolt NASA SOH Trust Engine API",
        "version": "2.1.2",
        "status": "operational",
        "environment": environment,
        "engine": "Standard (Random Forest + Scikit-Learn)",
        "endpoints": {
            "test": "/api/test",
            "assets": "/api/assets",
            "model_info": "/api/model-info",
            "visualizations": "/api/visualizations/{plot_name}"
        }
    }


@app.get("/api/test")
async def test_route():
    """Direct test route to verify sub-paths are reachable."""
    return {"message": "Sub-path connectivity verified!", "status": "200 OK"}


@app.get("/debug/routes")
async def list_routes():
    """Diagnostic endpoint to see what the server actually knows."""
    routes = []
    for route in app.routes:
        routes.append({
            "path": route.path,
            "name": route.name,
            "methods": list(route.methods) if hasattr(route, "methods") else []
        })
    return {
        "registered_count": len(routes),
        "registered_routes": routes, 
        "file_structure": os.listdir(".")
    }


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    model_exists = os.path.exists('./models/nasa_soh_model.pkl')
    pipeline_exists = os.path.exists('./models/feature_pipeline_causal.pkl')
    
    return {
        "status": "healthy" if (model_exists and pipeline_exists) else "degraded",
        "model_loaded": model_exists,
        "pipeline_loaded": pipeline_exists
    }


@app.post("/api/predict", response_model=HealthPredictionResponse)
async def predict_health_metrics(input_data: BMSDataInput):
    """
    Main endpoint: Predict battery health from BMS data.
    Uses NASA-trained model for Trust Engine proof.
    """
    try:
        # Convert to DataFrame
        data_dict = {
            'cycle_index': [d.cycle_index for d in input_data.data],
            'voltage': [d.voltage for d in input_data.data],
            'current': [d.current for d in input_data.data],
            'temperature': [d.temperature for d in input_data.data],
            'capacity': [d.capacity for d in input_data.data]
        }
        
        if input_data.data[0].timestamp:
            data_dict['timestamp'] = [d.timestamp for d in input_data.data]
        
        df = pd.DataFrame(data_dict)
        
        # Get predictions
        predictions = predict_soh(
            battery_data=df,
            model_path='./models/nasa_soh_model.pkl',
            pipeline_path='./models/feature_pipeline_causal.pkl',
            initial_capacity=input_data.initial_capacity,
            forecast_horizons=input_data.forecast_horizons
        )
        
        # Load model info
        model_info = get_model_info_internal()
        
        # Add model info to response
        predictions['model_info'] = {
            'model_type': model_info['model_type'],
            'training_data': model_info['training_data'],
            'validation_data': model_info['validation_data'],
            'r2_score': str(model_info['metrics']['r2']),
            'version': model_info['version']
        }
        
        return HealthPredictionResponse(**predictions)
        
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=f"Model files not found: {e}. Please train the model first.")
    except ValueError as e:
        raise HTTPException(status_code=400, detail=f"Invalid input data: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")


@app.post("/api/predict-csv")
async def predict_from_csv(file: UploadFile = File(...)):
    """
    Predict from CSV file upload.
    CSV should have columns: cycle_index, voltage, current, temperature, capacity
    """
    try:
        # Read CSV
        contents = await file.read()
        df = pd.read_csv(pd.io.common.BytesIO(contents))
        
        # Validate columns
        required_cols = ['cycle_index', 'voltage', 'current', 'temperature', 'capacity']
        missing_cols = [col for col in required_cols if col not in df.columns]
        if missing_cols:
            raise ValueError(f"Missing required columns: {missing_cols}")
        
        # Get predictions
        predictions = predict_soh(
            battery_data=df,
            model_path='./models/nasa_soh_model.pkl',
            pipeline_path='./models/feature_pipeline_causal.pkl'
        )
        
        return predictions
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"CSV processing failed: {str(e)}")


@app.get("/api/model-info", response_model=ModelInfoResponse)
async def get_model_info():
    """Get NASA model information for Trust Engine proof."""
    try:
        info = get_model_info_internal()
        return ModelInfoResponse(**info)
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="Model files not found. Please train the model first.")


def get_model_info_internal() -> Dict:
    """Internal function to get model info."""
    # Try to load expert model metadata first
    expert_model_path = './models/nasa_soh_model.pkl'
    
    if os.path.exists(expert_model_path):
        with open(expert_model_path, 'rb') as f:
            model_data = pickle.load(f)
            
        # Use expert model metadata if available
        if 'metrics' in model_data and 'val_r2' in model_data['metrics']:
            metrics = model_data['metrics']
            return {
                'model_type': model_data.get('model_type', 'RandomForestRegressor'),
                'training_data': ', '.join(model_data.get('training_batteries', ['B0005', 'B0006'])),
                'validation_data': model_data.get('validation_battery', 'B0007'),
                'metrics': {
                    'r2': float(metrics.get('val_r2', metrics.get('train_r2', 0.0))),
                    'mae': float(metrics.get('val_mae', metrics.get('train_mae', 0.0))),
                    'rmse': float(metrics.get('val_rmse', metrics.get('train_rmse', 0.0))),
                    'mape': float(metrics.get('val_mape', metrics.get('train_mape', 0.0)))
                },
                'feature_count': len(model_data.get('feature_cols', [])),
                'version': model_data.get('version', '2.0_expert'),
                'oob_score': metrics.get('oob_score')
            }
    
    # Fallback to model_comparison.csv
    comparison_df = pd.read_csv('./data/model_comparison.csv')
    best_model = comparison_df.iloc[0]
    
    # Load feature pipeline
    if os.path.exists('./models/feature_pipeline_causal.pkl'):
        with open('./models/feature_pipeline_causal.pkl', 'rb') as f:
            pipeline_info = pickle.load(f)
            feature_count = len(pipeline_info['feature_cols'])
            version = pipeline_info.get('version', '1.0_causal')
    else:
        feature_count = 47  # From expert model
        version = '1.0_causal'
    
    return {
        'model_type': 'NASA-trained Random Forest (Causal Features)',
        'training_data': 'NASA B0005, B0006 (168 cycles each)',
        'validation_data': 'NASA B0007 (168 cycles)',
        'metrics': {
            'r2': float(best_model['R² Score']),
            'mae': float(best_model['MAE']),
            'rmse': float(best_model['RMSE']),
            'mape': float(best_model['MAPE (%)'])
        },
        'feature_count': feature_count,
        'version': version
    }


@app.get("/api/visualizations/{plot_name}")
async def get_visualization(plot_name: str):
    """
    Serve visualization images for PDF reports.
    Returns base64-encoded image.
    """
    plot_map = {
        'health_forecast': 'health_forecast.png',
        'error_growth': 'error_growth.png',
        'model_simplicity': 'model_simplicity.png',
        'model_comparison': '01_model_comparison.png',
        'actual_vs_predicted': '02_actual_vs_predicted.png',
        'time_series': '03_time_series_prediction.png',
        'rul_prediction': '04_rul_prediction.png',
        'feature_importance': '05_feature_importance.png',
        'transfer_learning': '06_transfer_learning.png'
    }
    
    filename = plot_map.get(plot_name, plot_name)
    plot_path = f'./plots/{filename}'
    
    if not os.path.exists(plot_path):
        raise HTTPException(status_code=404, detail=f"Plot not found: {plot_name}")
    
    try:
        with open(plot_path, 'rb') as f:
            image_data = base64.b64encode(f.read()).decode()
        return {
            'plot_name': plot_name,
            'image': f'data:image/png;base64,{image_data}',
            'format': 'png'
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to load plot: {str(e)}")


@app.get("/api/all-visualizations")
async def get_all_visualizations():
    """Get all visualization plots for PDF reports."""
    plots = [
        ('health_forecast', 'health_forecast.png'),
        ('error_growth', 'error_growth.png'),
        ('model_simplicity', 'model_simplicity.png'),
        ('model_comparison', '01_model_comparison.png'),
        ('actual_vs_predicted', '02_actual_vs_predicted.png'),
        ('time_series', '03_time_series_prediction.png'),
        ('rul_prediction', '04_rul_prediction.png'),
        ('feature_importance', '05_feature_importance.png'),
        ('transfer_learning', '06_transfer_learning.png')
    ]
    
    result = {}
    for plot_name, filename in plots:
        plot_path = f'./plots/{filename}'
        if os.path.exists(plot_path):
            try:
                with open(plot_path, 'rb') as f:
                    image_data = base64.b64encode(f.read()).decode()
                    result[plot_name] = f'data:image/png;base64,{image_data}'
            except Exception as e:
                print(f"Warning: Could not load {plot_name}: {e}")
    
    return {
        'visualizations': result,
        'count': len(result)
    }


@app.get("/api/validation-report")
async def get_validation_report():
    """Get validation report data."""
    report_path = './data/nasa_model_validation.csv'
    
    if not os.path.exists(report_path):
        raise HTTPException(status_code=404, detail="Validation report not found. Please run production_pipeline.py first.")
    
    try:
        df = pd.read_csv(report_path)
        return {
            'data': df.to_dict(orient='records'),
            'summary': {
                'total_cycles': len(df),
                'train_cycles': len(df[df['dataset'] == 'train']),
                'val_cycles': len(df[df['dataset'] == 'validation']),
                'mean_mae': df['abs_error'].mean(),
                'mean_rmse': np.sqrt((df['error'] ** 2).mean()),
                'mean_mape': df['error_pct'].mean()
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to load validation report: {str(e)}")


if __name__ == "__main__":
    import uvicorn
    print("="*80)
    print("Starting ZenVolt NASA SOH Trust Engine API")
    print("="*80)
    print("\nEndpoints:")
    print("  POST /api/predict - Predict battery health from BMS data")
    print("  GET  /api/model-info - Get model information")
    print("  GET  /api/visualizations/{plot_name} - Get visualization images")
    print("  GET  /health - Health check")
    print("\nStarting server on http://localhost:8000")
    print("="*80)
    
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")