#!/usr/bin/env python3
# --------------------------------------------------------
# SEEM -- Segment Everything Everywhere All At Once
# Copyright (c) 2022 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Xueyan Zou (xueyan@cs.wisc.edu), Jianwei Yang (jianwyan@microsoft.com)
# --------------------------------------------------------

# Hugging Face Spaces Launcher
import os
import sys
import subprocess
import logging
import time
import shutil

# Configure logging
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("SEEM-HF")

def run_command(cmd, description=None):
    """Run a shell command and log its output"""
    if description:
        logger.info(f"Running: {description}")
    logger.info(f"Command: {cmd}")
    
    try:
        process = subprocess.Popen(
            cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            universal_newlines=True
        )
        
        # Stream and log output in real-time
        for line in process.stdout:
            line = line.rstrip()
            logger.info(line)
            
        process.wait()
        return process.returncode == 0
    except Exception as e:
        logger.error(f"Error executing command: {e}")
        return False

def install_dependencies():
    """Install required dependencies"""
    # Update package lists
    run_command("apt-get update", "Updating package lists")
    
    # Check if ffmpeg is installed
    logger.info("Checking for ffmpeg...")
    if not run_command("which ffmpeg", "Checking ffmpeg"):
        logger.info("Installing ffmpeg...")
        run_command("apt-get install -y ffmpeg", "Installing ffmpeg")
    
    # Install Python dependencies
    logger.info("Installing Python dependencies...")
    
    # Install requirements
    if os.path.exists("assets/requirements/requirements.txt"):
        run_command("pip install -r assets/requirements/requirements.txt", "Installing base requirements")
    
    # Install custom requirements
    if os.path.exists("assets/requirements/requirements_custom.txt"):
        run_command("pip install -r assets/requirements/requirements_custom.txt", "Installing custom requirements")
    
    # Install detectron2-xyz if not already installed
    try:
        import detectron2
        logger.info("detectron2 is already installed")
    except ImportError:
        logger.info("Installing detectron2-xyz...")
        run_command("pip install git+https://github.com/MaureenZOU/detectron2-xyz.git", "Installing detectron2-xyz")

def setup_environment():
    """Set up the necessary directories and environment"""
    # Make sure demo/seem directory exists
    if not os.path.exists("demo/seem"):
        logger.error("demo/seem directory not found!")
        return False
    
    # Create examples directory if it doesn't exist
    os.makedirs('demo/seem/examples', exist_ok=True)
    
    # Make sure configs directory exists
    if not os.path.exists("configs"):
        logger.error("configs directory not found!")
        # Create empty configs directory
        os.makedirs('configs/seem', exist_ok=True)
        logger.info("Created configs directory")
        
    # Create default config file if it doesn't exist
    config_path = "configs/seem/focall_unicl_lang_demo.yaml"
    if not os.path.exists(config_path):
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        with open(config_path, 'w') as f:
            f.write("""# Model config for demo
model:
  type: seem
  sem_seg_head:
    type: UniHead
    num_classes: 133
    embed_dim: 768
    learning_type: end2end
    lang_encoder: unicl
    dim_feedforward: 1024
    dropout: 0.1
    activation: relu
    nheads: 8
    dec_layers: 9
    enc_layers: 0
    text_encoder: "bert-base-uncased"
    mask_dim: 256
    task_switch:
      semantic: true
      instance: true
      panoptic: true
      refimg: true
      grounding: true
      audio: false
      audio_grounding: false
      video: true
""")
        logger.info(f"Created default config file at {config_path}")
    
    # Create constants.py in utils if it doesn't exist
    utils_dir = "utils"
    os.makedirs(utils_dir, exist_ok=True)
    constants_path = os.path.join(utils_dir, "constants.py")
    if not os.path.exists(constants_path):
        with open(constants_path, 'w') as f:
            f.write("""# COCO Panoptic Classes
COCO_PANOPTIC_CLASSES = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
    'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat',
    'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
    'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
    'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
    'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',
    'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator',
    'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
""")
        logger.info(f"Created constants.py at {constants_path}")
    
    return True

def download_model_weights():
    """Download model weights if they don't exist"""
    model_path = "seem_focall_v0.pt"
    if not os.path.exists(model_path):
        logger.info(f"Downloading model weights {model_path}...")
        run_command(
            f"wget https://huggingface.co/xdecoder/SEEM/resolve/main/{model_path}",
            f"Downloading {model_path}"
        )
        logger.info(f"Downloaded {model_path}")
    else:
        logger.info(f"Model weights {model_path} already exist")

def main():
    """Main entry point"""
    logger.info("Starting SEEM Hugging Face Space")
    
    # Install dependencies
    install_dependencies()
    
    # Setup environment
    if not setup_environment():
        return
    
    # Download model weights
    download_model_weights()
    
    # Prepare to run the actual app
    app_path = "demo/seem/app.py"
    if not os.path.exists(app_path):
        logger.error(f"Application file not found at {app_path}!")
        return
    
    logger.info(f"Running application from {app_path}")
    
    # Run the app
    cmd = f"python {app_path}"
    run_command(cmd, "Running SEEM demo app")

if __name__ == "__main__":
    main()