import spaces
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from ultralytics import YOLO
import cv2
import numpy as np
import requests
import os

app = FastAPI()

# Inisialisasi model YOLO
model = YOLO('yolov8n.pt')
URI_HP_SERVER = "https://visionx.my.id/api/upload"

# 1. Pisahkan fungsi inferensi GPU (Wajib synchronous 'def', bukan 'async def')
@spaces.GPU
def inferensi_yolo(img):
    return model(img)

# 2. Endpoint API FastAPI
@app.post("/api/predict")
def predict_and_forward(image: UploadFile = File(...), info: str = Form("Web Detection")):
    # Baca buffer gambar secara synchronous
    contents = image.file.read()
    nparr = np.frombuffer(contents, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    # Deteksi AI via GPU Zero
    results = inferensi_yolo(img)
    hasil = results[0]
    output_path = "temp_result.jpg"
    cv2.imwrite(output_path, hasil.plot())

    # Teruskan hasil ke server HP Termux
    status_hp = "Belum terkirim"
    try:
        with open(output_path, 'rb') as f:
            files = {'image': ('hasil.jpg', f, 'image/jpeg')}
            res = requests.post(URI_HP_SERVER, files=files, data={'info': info}, timeout=15)
            status_hp = "Sukses tersimpan di HP" if res.status_code == 200 else f"Gagal HP: {res.status_code}"
    except Exception as e:
        status_hp = f"Error HP: {str(e)}"

    return {
        "status": "success",
        "detections": len(hasil.boxes),
        "hp_storage_status": status_hp
    }

# 3. Render frontend VisionX dari folder dist
if os.path.exists("dist"):
    app.mount("/assets", StaticFiles(directory="dist/assets"), name="assets")

    @app.get("/{full_path:path}")
    def serve_spa(full_path: str):
        file_path = os.path.join("dist", full_path)
        if os.path.exists(file_path) and os.path.isfile(file_path):
            return FileResponse(file_path)
        return FileResponse("dist/index.html")

# 4. Runner port 7860 Hugging Face Spaces
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)