import cv2
import numpy as np
import requests
import torch
import firebase_admin
import os
from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
from ultralytics import YOLO
from firebase_admin import credentials, firestore

# --- 0. THE "FORCE TRUST" SECURITY OVERRIDE ---
# This stops the (y/N) prompt by forcing Torch Hub to trust all sub-repos
import torch.hub

# We redefine the internal check to always return True (TRUST EVERYTHING)
torch.hub.trust_repo = lambda *args, **kwargs: True

# Set environment variables for Hugging Face writable directories
os.environ['TORCH_HOME'] = '/tmp/torch_cache'
os.environ['YOLO_CONFIG_DIR'] = '/tmp/ultralytics_config'

# --- 1. INITIALIZE MODELS ---
app = FastAPI()

@app.get("/")
def home():
    return {"status": "Sahl Express AI is Online", "region": "Tunisia"}

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

print("🚀 Starting Sahl Express Engine...")

# Load YOLOv8 (2D Segmentation)
try:
    yolo_model = YOLO('best.pt') 
    print("✅ YOLOv8 Loaded")
except Exception as e:
    print(f"❌ YOLO Load Error: {e}")

# Load MiDaS (Depth Estimation)
try:
    print("📥 Loading MiDaS (Security Bypass Active)...")
    
    # Using 'trust_repo=True' alongside our override above
    midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", trust_repo=True)
    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms", trust_repo=True)
    
    midas.to(device)
    midas.eval()
    transform = midas_transforms.small_transform
    print("✅ MiDaS Loaded Successfully!")
except Exception as e:
    print(f"❌ MiDaS Load Failed: {e}")

# --- 2. FIREBASE SETUP ---
try:
    # Ensure serviceAccount.json is uploaded to your HF Space Files tab
    cred = credentials.Certificate("serviceAccount.json")
    firebase_admin.initialize_app(cred)
    db = firestore.client()
    print("✅ Firebase Connected")
except Exception as e:
    print(f"⚠️ Firebase Error: {e}")

# Tunisian Reference Constants (cm)
REFERENCE_SIZES = {
    'reference_card': 8.56,    # ID Card
    'reference_paper': 21.0,   # A4 Paper
    'reference_coin': 2.8      # 1 Dinar
}

class ImageRequest(BaseModel):
    image_url: str
    delivery_id: str

def get_depth_map(img):
    """ Converts an image to a relative depth map """
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    input_batch = transform(img_rgb).to(device)
    
    with torch.no_grad():
        prediction = midas(input_batch)
        prediction = torch.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size=img.shape[:2],
            mode="bicubic",
            align_corners=False,
        ).squeeze()
    
    return prediction.cpu().numpy()

def perform_3d_measurement(image_url: str, delivery_id: str):
    try:
        # Download Image from Cloudinary/URL
        resp = requests.get(image_url)
        img_array = np.asarray(bytearray(resp.content), dtype=np.uint8)
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

        # A. Run AI Models
        yolo_results = yolo_model.predict(source=img, conf=0.4)[0]
        depth_map = get_depth_map(img)
        
        pixel_cm_ratio = None
        pkg_mask = None
        pkg_w_px, pkg_h_px = None, None

        # 1. Calibration: Find the reference object (e.g., Tunisian ID card)
        for i, box in enumerate(yolo_results.boxes):
            label = yolo_results.names[int(box.cls[0])]
            if label in REFERENCE_SIZES:
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                pixel_cm_ratio = (x2 - x1) / REFERENCE_SIZES[label]
                break

        # 2. Identification: Find the Package and its mask
        for i, box in enumerate(yolo_results.boxes):
            label = yolo_results.names[int(box.cls[0])]
            if label == 'package' and yolo_results.masks is not None:
                pkg_mask = yolo_results.masks.xy[i]
                rect = cv2.minAreaRect(pkg_mask.astype(np.int32))
                (_, _), (w, h), _ = rect
                pkg_w_px, pkg_h_px = w, h
                break

        # 3. 3D Volume Calculation
        if pixel_cm_ratio and pkg_w_px is not None:
            # Create a mask to sample depth data
            mask_img = np.zeros(depth_map.shape, dtype=np.uint8)
            cv2.fillPoly(mask_img, [pkg_mask.astype(np.int32)], 1)
            
            pkg_depth_val = np.median(depth_map[mask_img == 1])
            
            # Ground depth (dilating the package mask to find the floor)
            kernel = np.ones((30,30), np.uint8)
            dilated = cv2.dilate(mask_img, kernel, iterations=2)
            ground_depth_val = np.median(depth_map[(dilated - mask_img) == 1])
            
            # Convert Relative Depth to Real CM
            # TUNING_CONSTANT: 0.5 is a baseline; adjust after testing with real packages
            TUNING_CONSTANT = 0.5 
            depth_delta = abs(ground_depth_val - pkg_depth_val)
            real_h = round((depth_delta / pixel_cm_ratio) * TUNING_CONSTANT, 1)
            
            # Final 2D Dimensions
            real_w = round(pkg_w_px / pixel_cm_ratio, 1)
            real_l = round(pkg_h_px / pixel_cm_ratio, 1)
            
            if real_h < 0.5: real_h = 1.0 # Minimum thickness
            volume = round(real_w * real_l * real_h, 2)

            # 4. Update Firebase with the measured volume
            db.collection("orders").document(delivery_id).update({
                "volume_cm3": volume,
                "dimensions": f"{real_l}x{real_w}x{real_h} cm",
                "status": "Measured_3D"
            })
            print(f"📦 Success: {delivery_id} | Vol: {volume}cm3")

    except Exception as e:
        print(f"❌ Measurement Error: {e}")

@app.post("/measure")
async def measure_endpoint(request: ImageRequest, background_tasks: BackgroundTasks):
    # This runs the heavy AI work in the background so the app doesn't freeze
    background_tasks.add_task(perform_3d_measurement, request.image_url, request.delivery_id)
    return {"status": "processing"}

if __name__ == "__main__":
    import uvicorn
    # Port 7860 is required for Hugging Face Spaces
    uvicorn.run(app, host="0.0.0.0", port=7860)