import os
import gc
import time
import shutil
import logging
import psutil
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse
from deepface import DeepFace

# Setup Logging
logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger("AIMonitor")

app = FastAPI()

# ----------------------------------------------------------------------------
# Helper Pengukur RAM
# ----------------------------------------------------------------------------
def get_mem():
    ram = psutil.virtual_memory()
    proc = psutil.Process(os.getpid())
    app_gb = proc.memory_info().rss / (1024 ** 3)
    free_gb = ram.available / (1024 ** 3)
    return app_gb, free_gb

# Path database lokal (bawaan lama)
DB_PATH = "/app/face_db"
if not os.path.exists(DB_PATH):
    os.makedirs(DB_PATH)


# ============================================================================
# MIDDLEWARE RAM TRACKER 
# ============================================================================
@app.middleware("http")
async def ram_tracker_middleware(request, call_next):
    # Abaikan logging jika hanya healthcheck root "/"
    if request.url.path == "/":
        return await call_next(request)

    start_time = time.time()
    ram_awal, free_awal = get_mem()

    # Catat log saat request mulai masuk
    print("\n" + "-" * 70)
    logger.info(f"[HIT MASUK] {request.method} {request.url.path}")
    logger.info(f"   |-- RAM Awal App   : {ram_awal:.2f} GB")
    logger.info(f"   |-- Sisa RAM Server: {free_awal:.2f} GB")
    logger.info(f"   |-- Status         : Sedang memproses AI...")

    # Jalankan kode AI
    response = await call_next(request)

    # Catat durasi dan bersihkan RAM sementara (Garbage Collector)
    duration_ms = (time.time() - start_time) * 1000
    ram_proses, _ = get_mem()
    gc.collect()  # Langsung lepaskan RAM sementara
    ram_akhir, free_akhir = get_mem()

    # Cetak alur naik dan turunnya RAM
    logger.info(f"[SELESAI] Status: {response.status_code} | Durasi: {duration_ms:.0f} ms")
    logger.info(
        f"[SIKLUS RAM] {ram_awal:.2f} GB --(naik proses)--> {ram_proses:.2f} GB --(turun cooldown)--> {ram_akhir:.2f} GB"
    )
    logger.info(f"[SISA RAM SERVER] {free_akhir:.2f} GB")
    print("-" * 70 + "\n")

    return response


# ============================================================================
# CODINGAN MODEL AI 
# ============================================================================
@app.get("/")
async def root():
    ram_app, ram_free = get_mem()
    return {
        "status": "AI Face Recognition System is Running",
        "app_ram_gb": round(ram_app, 2),
        "available_ram_gb": round(ram_free, 2)
    }

@app.post("/api/represent")
async def represent_face(
    image: UploadFile = File(...), 
    name: str = Form(None),
    detector_backend: str = Form("retinaface")
):
    """Endpoint stateless untuk mengekstrak vektor embedding wajah (512 dimensi untuk Facenet512)"""
    temp_path = f"temp_{int(time.time()*1000)}.jpg"
    try:
        if name:
            logger.info(f"   |-- Subjek: {name} (detector: {detector_backend})")

        # Simpan gambar secara temporer
        with open(temp_path, "wb") as buffer:
            shutil.copyfileobj(image.file, buffer)

        # Hitung representasi vektor wajah
        representations = DeepFace.represent(
            img_path=temp_path,
            model_name="Facenet512",
            detector_backend=detector_backend,
            enforce_detection=False
        )

        if len(representations) > 0:
            embedding = representations[0]["embedding"]
            return {"success": True, "embedding": embedding}
        
        return JSONResponse(status_code=400, content={"success": False, "error": "Wajah tidak dapat dideteksi."})

    except Exception as e:
        logger.error(f"   [ERROR] Error in represent: {str(e)}")
        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)

@app.post("/api/register")
async def register_face(name: str = Form(...), image: UploadFile = File(...)):
    try:
        clean_name = name.replace(" ", "_").lower()
        img_path = os.path.join(DB_PATH, f"{clean_name}.jpg")
        with open(img_path, "wb") as buffer:
            shutil.copyfileobj(image.file, buffer)

        for file in os.listdir(DB_PATH):
            if file.endswith(".pkl"):
                os.remove(os.path.join(DB_PATH, file))

        return {"success": True, "message": f"Wajah {name} berhasil didaftarkan."}
    except Exception as e:
        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})

@app.post("/api/recognize")
async def recognize_face(image: UploadFile = File(...)):
    temp_path = f"temp_scan_{int(time.time()*1000)}.jpg"
    try:
        with open(temp_path, "wb") as buffer:
            shutil.copyfileobj(image.file, buffer)

        has_images = any(f.lower().endswith(('.jpg', '.jpeg', '.png')) for f in os.listdir(DB_PATH))
        if not has_images:
            return {"success": False, "name": "Tidak Dikenali"}

        dfs = DeepFace.find(
            img_path=temp_path,
            db_path=DB_PATH,
            model_name="VGG-Face", 
            enforce_detection=False,
            silent=True
        )

        if len(dfs) > 0 and not dfs[0].empty:
            result_path = dfs[0]["identity"][0]
            name = os.path.basename(result_path).replace(".jpg", "").replace("_", " ").title()
            return {"success": True, "name": name}
        
        return {"success": False, "name": "Tidak Dikenali"}
    except Exception as e:
        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)

@app.post("/api/reset-ai")
async def reset_ai():
    try:
        for filename in os.listdir(DB_PATH):
            file_path = os.path.join(DB_PATH, filename)
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        return {"success": True, "message": "Database AI lokal dibersihkan."}
    except Exception as e:
        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})