import gradio as gr
import os
import sys
import logging
import time
import uuid
import atexit
from concurrent.futures import ThreadPoolExecutor
from typing import Union, List, Tuple, Dict, Any

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Import spaces for ZeroGPU support
try:
    import spaces   
except ImportError:
    # Fallback for local development
    def spaces(func):
        return func

# Import other dependencies
import subprocess
import cv2
import numpy as np
import threading
import tempfile
import shutil
import glob
import json
import base64
import struct
import zlib
import argparse
import socket
import gc
from pathlib import Path
from einops import rearrange
from tempfile import TemporaryDirectory
from http.server import SimpleHTTPRequestHandler
from socketserver import ThreadingTCPServer
import socketserver
import http.server
import torch 
from huggingface_hub import hf_hub_download

# Import custom modules with error handling
try:
    from app_3rd.sam_utils.inference import SamPredictor, get_sam_predictor, run_inference
    from app_3rd.spatrack_utils.infer_track import get_tracker_predictor, run_tracker, get_points_on_a_grid
except ImportError as e:
    logger.error(f"Failed to import custom modules: {e}")
    raise

MAX_FRAMES = 80

try:
    import vggt
except:
    subprocess.run(["pip", "install", "-e", "./models/vggt"], check=True)
    sys.path.append("/home/user/app/models/vggt")

# init the model
os.environ["VGGT_DIR"] = hf_hub_download("facebook/VGGT-1B", "model.pt")

if os.environ.get("VGGT_DIR", None) is not None:
    from vggt.models.vggt import VGGT
    from vggt.utils.load_fn import preprocess_image
    from vggt.utils.pose_enc import pose_encoding_to_extri_intri
    vggt_model = VGGT()
    vggt_model.load_state_dict(torch.load(os.environ.get("VGGT_DIR")))
    vggt_model.eval()
    vggt_model = vggt_model.to("cuda")

# Global model initialization
print("🚀 Initializing global models...")

def init_global_models():
    """Initialize global models (CPU only for ZeroGPU compatibility)"""
    try:
        print("🔧 Loading SAM predictor...")
        sam_predictor = get_sam_predictor()
        print("✅ SAM predictor loaded successfully")
        # Keep on CPU for ZeroGPU - will be moved to GPU in the decorated function

        print("🔧 Loading tracker models...")
        out_dir = os.path.join("temp_init", "results")
        os.makedirs(out_dir, exist_ok=True)
        tracker_model, tracker_viser = get_tracker_predictor(out_dir, vo_points=756)
        print("✅ Tracker models loaded successfully")
        # Keep on CPU for ZeroGPU - will be moved to GPU in the decorated function
        
        print("✅ All models initialized successfully!")
        return True
        
    except Exception as e:
        print(f"❌ Error initializing models: {e}")
        import traceback
        traceback.print_exc()
        return False

# Initialize models at startup

# Thread pool for delayed deletion
thread_pool_executor = ThreadPoolExecutor(max_workers=2)

def delete_later(path: Union[str, os.PathLike], delay: int = 600):
    """Delete file or directory after specified delay (default 10 minutes)"""
    def _delete():
        try:
            if os.path.isfile(path):
                os.remove(path)
            elif os.path.isdir(path):
                shutil.rmtree(path)
        except Exception as e:
            logger.warning(f"Failed to delete {path}: {e}")
    
    def _wait_and_delete():
        time.sleep(delay)
        _delete()
    
    thread_pool_executor.submit(_wait_and_delete)
    atexit.register(_delete)

def create_user_temp_dir():
    """Create a unique temporary directory for each user session"""
    session_id = str(uuid.uuid4())[:8]  # Short unique ID
    temp_dir = os.path.join("temp", f"session_{session_id}")
    os.makedirs(temp_dir, exist_ok=True)
    
    # Schedule deletion after 10 minutes
    delete_later(temp_dir, delay=600)
    
    return temp_dir

# Wrap the core GPU functions with @spaces.GPU
@spaces.GPU
def gpu_run_sam(image, points, boxes):
    """GPU-accelerated SAM inference"""
    # Initialize SAM predictor inside GPU function
    predictor = get_sam_predictor()
    
    # Ensure predictor is on GPU - handle different SAM predictor types
    try:
        if hasattr(predictor, 'model'):
            # For transformers SAM
            predictor.model = predictor.model.cuda()
        elif hasattr(predictor, 'sam'):
            # For segment-anything SAM
            predictor.sam = predictor.sam.cuda()
        elif hasattr(predictor, 'to'):
            # Generic PyTorch model
            predictor = predictor.to('cuda')
        
        # Also ensure image is on the right device if it's a tensor
        if hasattr(image, 'cuda'):
            image = image.cuda()
            
    except Exception as e:
        print(f"Warning: Could not move predictor to GPU: {e}")
    
    return run_inference(predictor, image, points, boxes)

@spaces.GPU
def gpu_run_tracker(temp_dir, video_name, grid_size, vo_points, fps):
    """GPU-accelerated tracking"""
    import torchvision.transforms as T
    import decord
    
    # Initialize tracker model inside GPU function
    out_dir = os.path.join(temp_dir, "results")
    os.makedirs(out_dir, exist_ok=True)
    tracker_model, tracker_viser = get_tracker_predictor(out_dir, vo_points=vo_points)
    
    # Setup paths
    video_path = os.path.join(temp_dir, f"{video_name}.mp4")
    mask_path = os.path.join(temp_dir, f"{video_name}.png")
    out_dir = os.path.join(temp_dir, "results")
    os.makedirs(out_dir, exist_ok=True)
    
    # Load video using decord
    video_reader = decord.VideoReader(video_path)
    video_tensor = torch.from_numpy(video_reader.get_batch(range(len(video_reader))).asnumpy()).permute(0, 3, 1, 2)  # Convert to tensor and permute to (N, C, H, W)
    # resize make sure the shortest side is 336
    h, w = video_tensor.shape[2:]
    scale = max(224 / h, 224 / w)
    if scale < 1:
        new_h, new_w = int(h * scale), int(w * scale)
        video_tensor = T.Resize((new_h, new_w))(video_tensor)
    video_tensor = video_tensor[::fps].float()[:MAX_FRAMES]
    
    # Move video tensor to GPU
    video_tensor = video_tensor.cuda()
    print(f"Video tensor shape: {video_tensor.shape}, device: {video_tensor.device}")
    
    depth_tensor = None
    intrs = None
    extrs = None
    data_npz_load = {}

    # run vggt 
    if os.environ.get("VGGT_DIR", None) is not None:
        # process the image tensor
        video_tensor = preprocess_image(video_tensor)[None]
        with torch.no_grad():
            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                # Predict attributes including cameras, depth maps, and point maps.
                aggregated_tokens_list, ps_idx = vggt_model.aggregator(video_tensor.cuda()/255)
                pose_enc = vggt_model.camera_head(aggregated_tokens_list)[-1]
                # Extrinsic and intrinsic matrices, following OpenCV convention (camera from world)
                extrinsic, intrinsic = pose_encoding_to_extri_intri(pose_enc, video_tensor.shape[-2:])
                # Predict Depth Maps
                depth_map, depth_conf = vggt_model.depth_head(aggregated_tokens_list, video_tensor.cuda()/255, ps_idx)
        depth_tensor = depth_map.squeeze().cpu().numpy()
        extrs = np.eye(4)[None].repeat(len(depth_tensor), axis=0)
        extrs[:, :3, :4] = extrinsic.squeeze().cpu().numpy()
        intrs = intrinsic.squeeze().cpu().numpy()
        video_tensor = video_tensor.squeeze()
        #NOTE: 20% of the depth is not reliable
        threshold = depth_conf.squeeze()[0].view(-1).quantile(0.6).item()
        unc_metric = depth_conf.squeeze().cpu().numpy() > threshold

    # Load and process mask
    if os.path.exists(mask_path):
        mask = cv2.imread(mask_path)
        mask = cv2.resize(mask, (video_tensor.shape[3], video_tensor.shape[2]))
        mask = mask.sum(axis=-1)>0
    else:
        mask = np.ones_like(video_tensor[0,0].cpu().numpy())>0
        grid_size = 10

    # Get frame dimensions and create grid points
    frame_H, frame_W = video_tensor.shape[2:]
    grid_pts = get_points_on_a_grid(grid_size, (frame_H, frame_W), device="cuda")  # Create on GPU
    
    # Sample mask values at grid points and filter out points where mask=0
    if os.path.exists(mask_path):
        grid_pts_int = grid_pts[0].long()
        mask_values = mask[grid_pts_int.cpu()[...,1], grid_pts_int.cpu()[...,0]]
        grid_pts = grid_pts[:, mask_values]
    
    query_xyt = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2)[0].cpu().numpy()

    print(f"Query points shape: {query_xyt.shape}")
    # Run model inference
    with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
        (
            c2w_traj, intrs, point_map, conf_depth,
            track3d_pred, track2d_pred, vis_pred, conf_pred, video
        ) = tracker_model.forward(video_tensor, depth=depth_tensor,
                            intrs=intrs, extrs=extrs, 
                            queries=query_xyt,
                            fps=1, full_point=False, iters_track=4,
                            query_no_BA=True, fixed_cam=False, stage=1,
                            support_frame=len(video_tensor)-1, replace_ratio=0.2) 
        
        # Resize results to avoid too large I/O Burden
        max_size = 224
        h, w = video.shape[2:]
        scale = min(max_size / h, max_size / w)
        if scale < 1:
            new_h, new_w = int(h * scale), int(w * scale)
            video = T.Resize((new_h, new_w))(video)
            video_tensor = T.Resize((new_h, new_w))(video_tensor)
            point_map = T.Resize((new_h, new_w))(point_map)
            track2d_pred[...,:2] = track2d_pred[...,:2] * scale
            intrs[:,:2,:] = intrs[:,:2,:] * scale
            conf_depth = T.Resize((new_h, new_w))(conf_depth)
        
        # Visualize tracks
        tracker_viser.visualize(video=video[None],
                        tracks=track2d_pred[None][...,:2],
                        visibility=vis_pred[None],filename="test")
                        
        # Save in tapip3d format
        data_npz_load["coords"] = (torch.einsum("tij,tnj->tni", c2w_traj[:,:3,:3], track3d_pred[:,:,:3].cpu()) + c2w_traj[:,:3,3][:,None,:]).numpy()
        data_npz_load["extrinsics"] = torch.inverse(c2w_traj).cpu().numpy()
        data_npz_load["intrinsics"] = intrs.cpu().numpy()
        data_npz_load["depths"] = point_map[:,2,...].cpu().numpy()
        data_npz_load["video"] = (video_tensor).cpu().numpy()/255
        data_npz_load["visibs"] = vis_pred.cpu().numpy()
        data_npz_load["confs"] = conf_pred.cpu().numpy()
        data_npz_load["confs_depth"] = conf_depth.cpu().numpy()
        np.savez(os.path.join(out_dir, f'result.npz'), **data_npz_load)
            
    return os.path.join(out_dir, "result.npz"), os.path.join(out_dir, "test_pred_track.mp4")

def compress_and_write(filename, header, blob):
    header_bytes = json.dumps(header).encode("utf-8")
    header_len = struct.pack("<I", len(header_bytes))
    with open(filename, "wb") as f:
        f.write(header_len)
        f.write(header_bytes)
        f.write(blob)

def process_point_cloud_data(npz_file, width=256, height=192, fps=4):
    fixed_size = (width, height)
    
    data = np.load(npz_file)
    extrinsics = data["extrinsics"]
    intrinsics = data["intrinsics"]
    trajs = data["coords"]
    T, C, H, W = data["video"].shape
    
    fx = intrinsics[0, 0, 0]
    fy = intrinsics[0, 1, 1]
    fov_y = 2 * np.arctan(H / (2 * fy)) * (180 / np.pi)
    fov_x = 2 * np.arctan(W / (2 * fx)) * (180 / np.pi)
    original_aspect_ratio = (W / fx) / (H / fy)
    
    rgb_video = (rearrange(data["video"], "T C H W -> T H W C") * 255).astype(np.uint8)
    rgb_video = np.stack([cv2.resize(frame, fixed_size, interpolation=cv2.INTER_AREA)
                          for frame in rgb_video])
    
    depth_video = data["depths"].astype(np.float32)
    if "confs_depth" in data.keys():
        confs = (data["confs_depth"].astype(np.float32) > 0.5).astype(np.float32)
        depth_video = depth_video * confs
    depth_video = np.stack([cv2.resize(frame, fixed_size, interpolation=cv2.INTER_NEAREST)
                            for frame in depth_video])
    
    scale_x = fixed_size[0] / W
    scale_y = fixed_size[1] / H
    intrinsics = intrinsics.copy()
    intrinsics[:, 0, :] *= scale_x
    intrinsics[:, 1, :] *= scale_y
    
    min_depth = float(depth_video.min()) * 0.8
    max_depth = float(depth_video.max()) * 1.5
    
    depth_normalized = (depth_video - min_depth) / (max_depth - min_depth)
    depth_int = (depth_normalized * ((1 << 16) - 1)).astype(np.uint16)
    
    depths_rgb = np.zeros((T, fixed_size[1], fixed_size[0], 3), dtype=np.uint8)
    depths_rgb[:, :, :, 0] = (depth_int & 0xFF).astype(np.uint8)
    depths_rgb[:, :, :, 1] = ((depth_int >> 8) & 0xFF).astype(np.uint8)
    
    first_frame_inv = np.linalg.inv(extrinsics[0])
    normalized_extrinsics = np.array([first_frame_inv @ ext for ext in extrinsics])
    
    normalized_trajs = np.zeros_like(trajs)
    for t in range(T):
        homogeneous_trajs = np.concatenate([trajs[t], np.ones((trajs.shape[1], 1))], axis=1)
        transformed_trajs = (first_frame_inv @ homogeneous_trajs.T).T
        normalized_trajs[t] = transformed_trajs[:, :3]
    
    arrays = {
        "rgb_video": rgb_video,
        "depths_rgb": depths_rgb,
        "intrinsics": intrinsics,
        "extrinsics": normalized_extrinsics,
        "inv_extrinsics": np.linalg.inv(normalized_extrinsics),
        "trajectories": normalized_trajs.astype(np.float32),
        "cameraZ": 0.0
    }
    
    header = {}
    blob_parts = []
    offset = 0
    for key, arr in arrays.items():
        arr = np.ascontiguousarray(arr)
        arr_bytes = arr.tobytes()
        header[key] = {
            "dtype": str(arr.dtype),
            "shape": arr.shape,
            "offset": offset,
            "length": len(arr_bytes)
        }
        blob_parts.append(arr_bytes)
        offset += len(arr_bytes)
    
    raw_blob = b"".join(blob_parts)
    compressed_blob = zlib.compress(raw_blob, level=9)
    
    header["meta"] = {
        "depthRange": [min_depth, max_depth],
        "totalFrames": int(T),
        "resolution": fixed_size,
        "baseFrameRate": fps,
        "numTrajectoryPoints": normalized_trajs.shape[1],
        "fov": float(fov_y),
        "fov_x": float(fov_x),
        "original_aspect_ratio": float(original_aspect_ratio),
        "fixed_aspect_ratio": float(fixed_size[0]/fixed_size[1])
    }
    
    # Create temporary file for compression
    temp_data_file = f'./temp_data_{int(time.time())}.bin'
    compress_and_write(temp_data_file, header, compressed_blob)
    
    # Read the compressed data and encode to base64
    with open(temp_data_file, "rb") as f:
        encoded_blob = base64.b64encode(f.read()).decode("ascii")
    
    # Clean up temporary file
    os.unlink(temp_data_file)
    
    # Read the HTML template and inject the base64 data
    with open('./_viz/viz_template.html') as f:
        html_template = f.read()
    
    # Inject the base64 data into the HTML
    html_content = html_template.replace(
        "<head>",
        f"<head>\n<script>window.embeddedBase64 = `{encoded_blob}`;</script>"
    )
    
    return html_content

def numpy_to_base64(arr):
    """Convert numpy array to base64 string"""
    return base64.b64encode(arr.tobytes()).decode('utf-8')

def base64_to_numpy(b64_str, shape, dtype):
    """Convert base64 string back to numpy array"""
    return np.frombuffer(base64.b64decode(b64_str), dtype=dtype).reshape(shape)

def get_video_name(video_path):
    """Extract video name without extension"""
    return os.path.splitext(os.path.basename(video_path))[0]

# Backend API Functions
def backend_upload_video(video_path: str) -> Dict[str, Any]:
    """Backend API for video upload"""
    try:
        # Create user-specific temporary directory
        user_temp_dir = create_user_temp_dir()
        
        # Get original video name
        video_name = get_video_name(video_path)
        temp_video_path = os.path.join(user_temp_dir, f"{video_name}.mp4")
        shutil.copy(video_path, temp_video_path)

        print(f"Video saved to: {temp_video_path}")
        
        # Extract first frame
        cap = cv2.VideoCapture(temp_video_path)
        success, frame = cap.read()
        cap.release()
        
        if not success:
            return {"success": False, "error": "Failed to read video"}
        
        # Resize frame to have minimum side length of 336
        h, w = frame.shape[:2]
        scale = 336 / min(h, w)
        new_h, new_w = int(h * scale)//2*2, int(w * scale)//2*2
        frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # Convert frame to base64 string for storage, include temp_dir info
        frame_data = {
            'data': numpy_to_base64(frame),
            'shape': frame.shape,
            'dtype': str(frame.dtype),
            'temp_dir': user_temp_dir  # Store temp directory path
        }
        
        return {
            "success": True,
            "original_image_state": json.dumps(frame_data),
            "display_image": frame,
            "selected_points": [],
            "temp_dir": user_temp_dir
        }
        
    except Exception as e:
        logger.error(f"Error in backend_upload_video: {e}")
        return {"success": False, "error": str(e)}

def backend_select_point(original_img: str, sel_pix: list, point_type: str, point_x: int, point_y: int) -> Dict[str, Any]:
    """Backend API for point selection"""
    try:
        # Convert stored image data back to numpy array
        frame_data = json.loads(original_img)
        original_img = base64_to_numpy(frame_data['data'], frame_data['shape'], frame_data['dtype'])
        temp_dir = frame_data.get('temp_dir', 'temp')  # Get user-specific temp dir
            
        # Create a display image for visualization
        display_img = original_img.copy()
        # Create a new list instead of modifying the existing one
        new_sel_pix = sel_pix.copy() if sel_pix else []
        new_sel_pix.append(((point_x, point_y), 1 if point_type == 'positive_point' else 0))
        
        # Run SAM inference
        o_masks = gpu_run_sam(original_img, new_sel_pix, [])
        
        # Draw points on display image
        COLORS = [(0, 0, 255), (0, 255, 255)]  # BGR: Red for negative, Yellow for positive
        MARKERS = [1, 5]  # Cross for negative, Star for positive
        MARKER_SIZE = 8  # Increased marker size
        
        for point, label in new_sel_pix:
            cv2.drawMarker(display_img, point, COLORS[label], markerType=MARKERS[label], markerSize=MARKER_SIZE, thickness=2)
        
        # Draw mask overlay on display image
        if o_masks:
            # Get the final mask (which is already processed as pos_mask - neg_mask)
            mask = o_masks[0][0]  # Get first mask

            # Create a light blue overlay
            overlay = display_img.copy()
            overlay[mask.squeeze()!=0] = [20, 60, 200]  # Light blue in BGR
            # Blend with original image with lower alpha
            display_img = cv2.addWeighted(overlay, 0.6, display_img, 0.4, 0)
        
        # Save mask
        if o_masks:
            video_files = glob.glob(os.path.join(temp_dir, "*.mp4"))
            if video_files:
                video_name = get_video_name(video_files[0])
                for mask, _ in o_masks:
                    o_mask = np.uint8(mask.squeeze() * 255)
                    o_file = os.path.join(temp_dir, f"{video_name}.png")
                    cv2.imwrite(o_file, o_mask)
        
        return {
            "success": True,
            "display_image": display_img,
            "selected_points": new_sel_pix
        }
        
    except Exception as e:
        logger.error(f"Error in backend_select_point: {e}")
        return {"success": False, "error": str(e)}

def backend_reset_points(original_img: str, sel_pix: list) -> Dict[str, Any]:
    """Backend API for resetting points"""
    try:
        # Convert stored image data back to numpy array
        frame_data = json.loads(original_img)
        original_img = base64_to_numpy(frame_data['data'], frame_data['shape'], frame_data['dtype'])
        temp_dir = frame_data.get('temp_dir', 'temp')  # Get user-specific temp dir
            
        # Create a display image for visualization (just the original image)
        display_img = original_img.copy()
        
        # Clear all points
        new_sel_pix = []
        
        # Clear any existing masks in user's temp directory
        for mask_file in glob.glob(os.path.join(temp_dir, "*.png")):
            try:
                os.remove(mask_file)
            except Exception as e:
                logger.warning(f"Failed to remove mask file {mask_file}: {e}")
        
        return {
            "success": True,
            "display_image": display_img,
            "selected_points": new_sel_pix
        }
        
    except Exception as e:
        logger.error(f"Error in backend_reset_points: {e}")
        return {"success": False, "error": str(e)}

def backend_run_tracker(grid_size: int, vo_points: int, fps: int, original_image_state: str) -> Dict[str, Any]:
    """Backend API for running tracker and visualization"""
    try:
        # Get user's temp directory from stored frame data
        frame_data = json.loads(original_image_state)
        temp_dir = frame_data.get('temp_dir', 'temp')
        
        video_files = glob.glob(os.path.join(temp_dir, "*.mp4"))
        if not video_files:
            return {"success": False, "error": "No video file found"}
        
        video_path = video_files[0]
        video_name = get_video_name(video_path)
        
        # Run tracker
        npz_path, track2d_video = gpu_run_tracker(temp_dir, video_name, grid_size, vo_points, fps)
        
        # Generate HTML content
        html_content = process_point_cloud_data(npz_path)
        
        # Schedule deletion of generated files
        if os.path.exists(track2d_video):
            delete_later(track2d_video, delay=600)
        if os.path.exists(npz_path):
            delete_later(npz_path, delay=600)
        
        return {
            "success": True,
            "viz_html": html_content,
            "track_video_path": track2d_video
        }
        
    except Exception as e:
        logger.error(f"Error in backend_run_tracker: {e}")
        return {"success": False, "error": str(e)}

# Remove the separate interfaces and create a unified API handler
def unified_api_handler(function_type: str, *args) -> Dict[str, Any]:
    """Unified API handler for all backend functions"""
    try:
        if function_type == "upload_video":
            # args[0] should be the video file
            return backend_upload_video(args[0])
        elif function_type == "select_point":
            # args: original_img, sel_pix, point_type, point_x, point_y
            return backend_select_point(args[0], args[1], args[2], args[3], args[4])
        elif function_type == "reset_points":
            # args: original_img, sel_pix
            return backend_reset_points(args[0], args[1])
        elif function_type == "run_tracker":
            # args: grid_size, vo_points, fps, original_image_state
            return backend_run_tracker(args[0], args[1], args[2], args[3])
        else:
            return {"success": False, "error": f"Unknown function type: {function_type}"}
    except Exception as e:
        logger.error(f"Error in unified_api_handler: {e}")
        return {"success": False, "error": str(e)}

# Create the main unified API interface
main_api = gr.Interface(
    fn=unified_api_handler,
    inputs=[
        gr.Dropdown(
            choices=["upload_video", "select_point", "reset_points", "run_tracker"],
            label="Function Type",
            value="upload_video"
        ),
        gr.File(label="Video File (for upload_video)", file_types=[".mp4", ".avi", ".mov"]),
        gr.Textbox(label="Original Image State", value=""),
        gr.JSON(label="Selected Points", value=[]),
        gr.Radio(choices=['positive_point', 'negative_point'], label="Point Type", value='positive_point'),
        gr.Number(label="Point X", value=0),
        gr.Number(label="Point Y", value=0),
        gr.Number(label="Grid Size", value=50),
        gr.Number(label="VO Points", value=756),
        gr.Number(label="FPS", value=3)
    ],
    outputs=[
        gr.JSON(label="Result")
    ],
    title="SpaTrackV2 Backend API",
    description="Unified Backend API for SpaTrackV2. This is a private Space that provides core functionality.",
    api_name="unified_api"
)

# Create additional interfaces for individual API functions for manual testing
select_point_api = gr.Interface(
    fn=backend_select_point,
    inputs=[
        gr.Textbox(label="Original Image State"),
        gr.JSON(label="Selected Points"),
        gr.Radio(choices=['positive_point', 'negative_point'], label="Point Type"),
        gr.Number(label="Point X"),
        gr.Number(label="Point Y")
    ],
    outputs=[
        gr.JSON(label="Result")
    ],
    title="Select Point API",
    description="API for selecting points on video frames"
)

reset_points_api = gr.Interface(
    fn=backend_reset_points,
    inputs=[
        gr.Textbox(label="Original Image State"),
        gr.JSON(label="Selected Points")
    ],
    outputs=[
        gr.JSON(label="Result")
    ],
    title="Reset Points API",
    description="API for resetting points"
)

tracker_api = gr.Interface(
    fn=backend_run_tracker,
    inputs=[
        gr.Number(label="Grid Size", value=50),
        gr.Number(label="VO Points", value=756),
        gr.Number(label="FPS", value=3),
        gr.Textbox(label="Original Image State")
    ],
    outputs=[
        gr.JSON(label="Result")
    ],
    title="Run Tracker API",
    description="API for running the tracking algorithm"
)

# Create a combined interface with tabs for manual testing
with gr.Blocks(title="SpaTrackV2 Backend API") as backend_app:
    gr.Markdown("# 🚀 SpaTrackV2 Backend API")
    gr.Markdown("This is a private backend Space that provides core SpaTrackV2 functionality.")
    
    with gr.Tabs():
        with gr.TabItem("Unified API"):
            main_api.render()
        
        with gr.TabItem("Upload Video"):
            upload_api = gr.Interface(
                fn=backend_upload_video,
                inputs=[gr.File(label="Upload Video", file_types=[".mp4", ".avi", ".mov"])],
                outputs=[gr.JSON(label="Result")],
                title="Upload Video API"
            )
            upload_api.render()
        
        with gr.TabItem("Select Point"):
            select_point_api.render()
        
        with gr.TabItem("Reset Points"):
            reset_points_api.render()
        
        with gr.TabItem("Run Tracker"):
            tracker_api.render()
        
        with gr.TabItem("API Info"):
            gr.Markdown("""
            ## Available API Functions
            
            ### Unified API
            - **Function**: `unified_api_handler`
            - **Input**: Function type + parameters
            - **Output**: JSON result
            
            ### Individual Functions
            
            #### 1. Upload Video
            - **Function**: `backend_upload_video`
            - **Input**: Video file
            - **Output**: Initial state and settings
            
            #### 2. Select Point
            - **Function**: `backend_select_point`
            - **Input**: Image state + point coordinates
            - **Output**: Updated image and points
            
            #### 3. Reset Points
            - **Function**: `backend_reset_points`
            - **Input**: Image state + points
            - **Output**: Reset image and empty points
            
            #### 4. Run Tracker
            - **Function**: `backend_run_tracker`
            - **Input**: Parameters + image state
            - **Output**: Visualization and tracking results
            
            ### 5. GPU Functions
            - `gpu_run_sam(image, points, boxes)`: GPU-accelerated SAM inference
            - `gpu_run_tracker(temp_dir, video_name, grid_size, vo_points, fps)`: GPU-accelerated tracking
            """)

if __name__ == "__main__":
    # Print startup information
    print("🚀 Starting SpaTrackV2 Backend Space...")
    print(f"🔧 Python version: {sys.version}")
    print(f"🔧 Working directory: {os.getcwd()}")
    print(f"🔧 GPU available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"🔧 GPU device: {torch.cuda.get_device_name(0)}")
        print(f"🔧 GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    
    print(f"🔧 Initializing models and GPU resources...")
    
    # Initialize global models
    init_success = init_global_models()
    if init_success:
        print("✅ Backend initialization complete!")
    else:
        print("❌ Backend initialization failed! Continuing with limited functionality...")
    
    print("📡 Starting Gradio backend interface...")
    print(f"🔧 Available GPU functions: {[name for name in globals() if name.startswith('gpu_')]}")
    
    # Launch the complete backend app (not just main_api)
    backend_app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,  # Backend shouldn't need sharing
        debug=True,
        show_error=True
    )