import os
import glob
import sys
import pandas as pd
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# Add src and subpackages to path
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(BASE_DIR, 'src'))

from rag import LLMExplainer, LogRetriever

app = FastAPI(title="Predictive Maintenance API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

PROCESSED_DIR = os.path.join(BASE_DIR, 'data', 'processed')
REPORTS_DIR = os.path.join(BASE_DIR, 'reports')
PLOTS_DIR = os.path.join(BASE_DIR, 'plots')
STATIC_DIR = os.path.join(BASE_DIR, 'static')
MODELS_DIR = os.path.join(BASE_DIR, 'models')
RAG_MODELS_DIR = os.path.join(BASE_DIR, 'models', 'rag')

os.makedirs(STATIC_DIR, exist_ok=True)
os.makedirs(REPORTS_DIR, exist_ok=True)
os.makedirs(PLOTS_DIR, exist_ok=True)
os.makedirs(PROCESSED_DIR, exist_ok=True)

# Initialize AI & RAG components
explainer = None
retriever = None

try:
    explainer = LLMExplainer(force_fallback=False)
    retriever = LogRetriever(models_dir=RAG_MODELS_DIR if os.path.exists(RAG_MODELS_DIR) else MODELS_DIR)
except Exception as e:
    print(f"Warning: Primary LLM initialization encountered exception: {e}")
    try:
        explainer = LLMExplainer(force_fallback=True)
        retriever = LogRetriever(models_dir=RAG_MODELS_DIR if os.path.exists(RAG_MODELS_DIR) else MODELS_DIR)
    except Exception as e2:
        print(f"Fallback initialization error: {e2}")


@app.get("/api/schedule")
def get_schedule():
    """Returns the adjusted production schedule."""
    csv_path = os.path.join(PROCESSED_DIR, 'adjusted_schedule.csv')
    if not os.path.exists(csv_path):
        return {"error": "Schedule not found. Run Phase 2 first.", "data": []}
    try:
        df = pd.read_csv(csv_path)
        return {"data": df.to_dict(orient="records")}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/reports")
def get_reports():
    """Returns list of generated anomaly reports."""
    report_files = glob.glob(os.path.join(REPORTS_DIR, 'machine_*_explanation_report.md'))
    reports = []
    for file_path in report_files:
        basename = os.path.basename(file_path)
        machine_id = basename.split('_')[1]
        reports.append({
            "machine_id": machine_id,
            "filename": basename
        })
    return {"reports": reports}


@app.get("/api/report/{filename}")
def get_report_content(filename: str):
    """Returns markdown content of an explanation report."""
    safe_name = os.path.basename(filename)
    file_path = os.path.join(REPORTS_DIR, safe_name)
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="Report not found")
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    return {"content": content}


@app.get("/api/status")
def get_overall_status():
    """Calculates overall factory health metrics."""
    csv_path = os.path.join(PROCESSED_DIR, 'adjusted_schedule.csv')
    if not os.path.exists(csv_path):
        return {"total_jobs": 0, "at_risk_machines": 0, "reassigned_jobs": 0, "healthy_jobs": 0}
        
    df = pd.read_csv(csv_path)
    total = len(df)
    reassigned = len(df[df['status'] == 'Rescheduled'])
    healthy = total - reassigned
    at_risk = reassigned
    
    return {
        "total_jobs": total,
        "at_risk_machines": at_risk,
        "reassigned_jobs": reassigned,
        "healthy_jobs": healthy
    }


# ---------------------------------------------------------------------------
#  NEW ENDPOINT 1: Fleet-wide RUL & Health Status for all 100 machines
# ---------------------------------------------------------------------------
@app.get("/api/fleet/rul")
def get_fleet_rul():
    """Returns RUL estimates and health status for every machine in the Azure fleet."""
    machines_path = os.path.join(BASE_DIR, 'data', 'raw', 'PdM_machines.csv')
    telemetry_path = os.path.join(BASE_DIR, 'data', 'raw', 'PdM_telemetry.csv')

    if not os.path.exists(machines_path):
        raise HTTPException(status_code=404, detail="PdM_machines.csv not found in data/raw/")
    if not os.path.exists(telemetry_path):
        raise HTTPException(status_code=404, detail="PdM_telemetry.csv not found in data/raw/")

    machines_df = pd.read_csv(machines_path)
    # Read only the last chunk of telemetry to get the latest readings per machine
    telemetry_df = pd.read_csv(telemetry_path)
    latest_telemetry = telemetry_df.sort_values('datetime').groupby('machineID').last().reset_index()

    fleet = []
    for _, mrow in machines_df.iterrows():
        mid = int(mrow['machineID'])
        model = mrow['model']
        age = int(mrow['age'])

        # Find latest sensor readings for this machine
        t_row = latest_telemetry[latest_telemetry['machineID'] == mid]
        if t_row.empty:
            fleet.append({
                "machine_id": mid, "model": model, "age": age,
                "rul_hours": 100.0, "failure_probability": 0.02,
                "status": "Healthy",
                "sensors": {"volt": None, "rotate": None, "pressure": None, "vibration": None}
            })
            continue

        volt = float(t_row['volt'].iloc[0])
        rotate = float(t_row['rotate'].iloc[0])
        pressure = float(t_row['pressure'].iloc[0])
        vibration = float(t_row['vibration'].iloc[0])

        # Same heuristic risk scoring as /api/predict
        risk_score = 0.0
        if volt > 180 or volt < 150: risk_score += 0.3
        if rotate > 500 or rotate < 350: risk_score += 0.3
        if pressure > 110 or pressure < 90: risk_score += 0.3
        if vibration > 50: risk_score += 0.4

        prob = min(max(risk_score, 0.02), 0.98)
        rul = max(100.0 - (prob * 100), 2.0)

        if prob > 0.40:
            status = "Critical"
        elif prob > 0.15:
            status = "At Risk"
        else:
            status = "Healthy"

        fleet.append({
            "machine_id": mid,
            "model": model,
            "age": age,
            "rul_hours": round(rul, 1),
            "failure_probability": round(prob, 4),
            "status": status,
            "sensors": {
                "volt": round(volt, 2),
                "rotate": round(rotate, 2),
                "pressure": round(pressure, 2),
                "vibration": round(vibration, 2)
            }
        })

    # Sort: Critical first, then At Risk, then Healthy
    status_order = {"Critical": 0, "At Risk": 1, "Healthy": 2}
    fleet.sort(key=lambda x: (status_order.get(x["status"], 3), x["rul_hours"]))

    summary = {
        "total_machines": len(fleet),
        "critical": sum(1 for m in fleet if m["status"] == "Critical"),
        "at_risk": sum(1 for m in fleet if m["status"] == "At Risk"),
        "healthy": sum(1 for m in fleet if m["status"] == "Healthy"),
    }

    return {"summary": summary, "fleet": fleet}


# ---------------------------------------------------------------------------
#  NEW ENDPOINT 2: Dynamic live rescheduling triggered from the dashboard
# ---------------------------------------------------------------------------
class RescheduleRequest(BaseModel):
    """Accepts per-machine predictions to trigger schedule adjustment."""
    predictions: dict  # { "3": {"failure_probability": 0.85, "rul": 12.0}, ... }
    failure_threshold: float = 0.70
    min_rul_hours: float = 48.0


@app.post("/api/reschedule")
def reschedule_live(req: RescheduleRequest):
    """Dynamically re-runs the ProductionScheduler with provided predictions."""
    from utils.scheduler import ProductionScheduler

    schedule_path = os.path.join(BASE_DIR, 'data', 'raw', 'production_schedule.csv')
    machines_path = os.path.join(BASE_DIR, 'data', 'raw', 'PdM_machines.csv')
    adjusted_path = os.path.join(PROCESSED_DIR, 'adjusted_schedule.csv')

    if not os.path.exists(schedule_path):
        raise HTTPException(status_code=404, detail="production_schedule.csv not found")
    if not os.path.exists(machines_path):
        raise HTTPException(status_code=404, detail="PdM_machines.csv not found")

    # Convert string keys to int keys (JSON keys are always strings)
    int_predictions = {}
    for k, v in req.predictions.items():
        int_predictions[int(k)] = {
            'failure_probability': float(v.get('failure_probability', 0.0)),
            'rul': float(v.get('rul', 500.0))
        }

    try:
        scheduler = ProductionScheduler(schedule_path, machines_path)
        adjusted_df = scheduler.adjust_schedule(
            predictions=int_predictions,
            output_adjusted_path=adjusted_path,
            failure_threshold=req.failure_threshold,
            min_rul_hours=req.min_rul_hours
        )

        total = len(adjusted_df)
        rescheduled = len(adjusted_df[adjusted_df['status'] == 'Rescheduled'])
        risk_alerts = len(adjusted_df[adjusted_df['status'] == 'Risk Alert'])

        return {
            "success": True,
            "total_jobs": total,
            "rescheduled_count": rescheduled,
            "risk_alert_count": risk_alerts,
            "healthy_count": total - rescheduled - risk_alerts,
            "schedule": adjusted_df.to_dict(orient="records")
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Rescheduling failed: {str(e)}")


class PredictRequest(BaseModel):
    machineID: int
    age: float
    volt: float
    rotate: float
    pressure: float
    vibration: float


@app.post("/api/predict")
def predict_live(req: PredictRequest):
    """Executes live prediction and generates an LLM diagnostic report."""
    if not explainer or not retriever:
        raise HTTPException(status_code=500, detail="Inference engine not initialized.")
        
    # Heuristic scoring simulation
    risk_score = 0.0
    if req.volt > 180 or req.volt < 150: risk_score += 0.3
    if req.rotate > 500 or req.rotate < 350: risk_score += 0.3
    if req.pressure > 110 or req.pressure < 90: risk_score += 0.3
    if req.vibration > 50: risk_score += 0.4
    
    prob = min(max(risk_score, 0.02), 0.98)
    rul = max(100.0 - (prob * 100), 2.0)
    
    # Risk Status Thresholds:
    # 0-15% -> Healthy, 16-40% -> At Risk, 41%+ -> Critical
    if prob > 0.40:
        status = "Critical"
    elif prob > 0.15:
        status = "At Risk"
    else:
        status = "Healthy"
    
    # Execute FAISS semantic query
    query_str = f"Machine ID: {req.machineID} telemetry anomalies volt: {req.volt} vib: {req.vibration}"
    hits = retriever.query(query_str, k=2)
    
    sensor_dict = {
        'volt': req.volt,
        'rotate': req.rotate,
        'pressure': req.pressure,
        'vibration': req.vibration
    }
    
    report = explainer.generate_explanation(
        machine_id=req.machineID,
        failure_prob=prob,
        rul=rul,
        historical_logs=hits,
        sensor_values=sensor_dict
    )
    
    return {
        "failure_probability": prob,
        "rul": rul,
        "status": status,
        "report": report
    }


# Mount static assets
app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="static")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="127.0.0.1", port=8000, reload=True)