import io
import os  # 🟢 Thêm dòng này để tránh lỗi NameError
import sys
import time
import numpy as np
import cv2
import torch
import torchvision
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import uvicorn

app = FastAPI()

# 🟢 Clone FastDepth nếu chưa có
fastdepth_path = "FastDepth"
if not os.path.exists(fastdepth_path):
    os.system("git clone https://github.com/dwofk/fast-depth.git FastDepth")

# 🟢 Thêm FastDepth vào sys.path để import được
sys.path.append(fastdepth_path)

# 🟢 Import FastDepth sau khi đã tải về
from FastDepth.models import MobileNetSkipAdd  

# 🟢 Load mô hình FastDepth
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MobileNetSkipAdd(output_size=(224, 224))  # 🟢 FastDepth hỗ trợ đầu ra 224x224
model.load_state_dict(torch.load(f"{fastdepth_path}/models/fastdepth_nyu.pt", map_location=device))
model.eval().to(device)

@app.post("/analyze_path/")
async def analyze_path(file: UploadFile = File(...)):
    # 🟢 Đọc file ảnh từ ESP32
    image_bytes = await file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    
    # 🟢 Chuyển đổi ảnh thành tensor phù hợp với FastDepth
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((224, 224)),  # FastDepth yêu cầu ảnh 224x224
        torchvision.transforms.ToTensor(),
    ])
    img_tensor = transform(image).unsqueeze(0).to(device)

    # 🟢 Bắt đầu đo thời gian dự đoán Depth Map
    start_time = time.time()
    
    # 🟢 Dự đoán Depth Map với FastDepth
    with torch.no_grad():
        depth_map = model(img_tensor).squeeze().cpu().numpy()

    end_time = time.time()
    print(f"⏳ FastDepth xử lý trong {end_time - start_time:.4f} giây")

    # 🟢 Kiểm tra kích thước Depth Map
    print(f"📏 Depth Map Shape: {depth_map.shape}")

    # 🟢 Đo thời gian xử lý đường đi
    start_detect_time = time.time()
    command = detect_path(depth_map)
    end_detect_time = time.time()
    print(f"⏳ detect_path() xử lý trong {end_detect_time - start_detect_time:.4f} giây")

    return {"command": command}

def detect_path(depth_map):
    """Phân tích đường đi từ ảnh Depth Map"""
    if len(depth_map.shape) != 2:  # 🟢 Kiểm tra nếu depth_map không phải 2D
        raise ValueError("Depth map không phải ảnh 2D hợp lệ!")

    h, w = depth_map.shape
    center_x = w // 2
    scan_y = h - 20  # Quét dòng gần đáy ảnh

    left_region = np.mean(depth_map[scan_y, :center_x])
    right_region = np.mean(depth_map[scan_y, center_x:])
    center_region = np.mean(depth_map[scan_y, center_x - 20:center_x + 20])

    if center_region > 200:
        return "forward"
    elif left_region > right_region:
        return "left"
    elif right_region > left_region:
        return "right"
    else:
        return "backward"

# 🟢 Chạy server FastAPI
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)