from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
import joblib
import numpy as np
import os
import threading
import time
from features import extract_features
import spaces  

app = FastAPI(title="Malicious URL Detector API")

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

models = {}

try:
    models["decision_tree"] = joblib.load('models/decision_tree_model.joblib')
    print("✅ Decision Tree loaded")
except Exception as e:
    print(f"❌ Decision Tree error: {e}")

# Load Random Forest
try:
    models["random_forest"] = joblib.load('models/random_forest_model.joblib')
    print("✅ Random Forest loaded")
except Exception as e:
    print(f"❌ Random Forest error: {e}")

# Load KNeighbors
try:
    for file in os.listdir('models/'):
        if 'kneighbor' in file.lower() or 'neighbor' in file.lower():
            models["kneighbors"] = joblib.load(f'models/{file}')
            print(f"✅ KNeighbors loaded from: {file}")
            break
except Exception as e:
    print(f"❌ KNeighbors error: {e}")

# Jika tidak ada model, buat dummy
if not models:
    print("⚠️ No models found, using dummy")
    from sklearn.tree import DecisionTreeClassifier
    models["decision_tree"] = DecisionTreeClassifier()
    X_dummy = np.random.rand(10, 22)
    y_dummy = np.random.randint(0, 4, 10)
    models["decision_tree"].fit(X_dummy, y_dummy)

class URLRequest(BaseModel):
    url: str
    model_type: str = "decision_tree"

CATEGORY_MAP = {
    0: "Aman (Benign)",
    1: "Defacement",
    2: "Phishing",
    3: "Malware"
}

@spaces.GPU  
@app.post("/analyze")
async def analyze_url(request: URLRequest):
    if not request.url:
        raise HTTPException(status_code=400, detail="URL tidak boleh kosong")
    
    if request.model_type not in models:
        request.model_type = list(models.keys())[0]
        
    active_model = models[request.model_type]
    
    try:
        input_data = extract_features(request.url)
        prediction = active_model.predict(input_data)[0]
        probabilities = active_model.predict_proba(input_data)[0]
        confidence = float(np.max(probabilities) * 100)
        
        return {
            "url": request.url,
            "used_model": request.model_type,
            "category_id": int(prediction),
            "category_name": CATEGORY_MAP.get(int(prediction), "Unknown"),
            "confidence": round(confidence, 2)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def read_root():
    return {"message": "Malicious URL Detector API is running", "status": "ok", "models": list(models.keys())}

@app.get("/health")
async def health():
    return {"status": "healthy", "models": list(models.keys())}

if os.path.exists("frontend"):
    app.mount("/css", StaticFiles(directory="frontend/css"), name="css")
    app.mount("/js", StaticFiles(directory="frontend/js"), name="js")
    app.mount("/assets", StaticFiles(directory="frontend/assets"), name="assets")

    @app.get("/")
    async def read_index():
        return FileResponse("frontend/index.html")
else:
    if os.path.exists("css"):
        app.mount("/css", StaticFiles(directory="css"), name="css")
    if os.path.exists("js"):
        app.mount("/js", StaticFiles(directory="js"), name="js")
    if os.path.exists("assets"):
        app.mount("/assets", StaticFiles(directory="assets"), name="assets")

    @app.get("/")
    async def read_index():
        if os.path.exists("index.html"):
            return FileResponse("index.html")
        return {"message": "Frontend not found"}

def keep_alive():
    while True:
        time.sleep(60)
        print("🔄 App is alive and running...")

import uvicorn

thread = threading.Thread(target=keep_alive, daemon=True)
thread.start()

uvicorn.run(app, host="0.0.0.0", port=7860)