import subprocess
import os
from pathlib import Path
import glob
import logging

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

def run_inference(image_path: str, output_dir: str):
    """
    Run Stable Fast 3D inference on the input image
    
    Args:
        image_path: Path to the input image
        output_dir: Directory to save output files
        
    Returns:
        Path to the generated GLB file
    """
    # Get the repository path
    repo_path = Path("/app/stable-fast-3d")

    
    logger.info(f"Running inference on image: {image_path}")
    logger.info(f"Output directory: {output_dir}")
    
    # Run the model
    command = [
        "python", "run.py",
        image_path,
        "--output-dir", output_dir,
        "--pretrained-model", "stabilityai/stable-fast-3d",
        "--texture-resolution", "2048",
        "--remesh_option", "quad"
    ]
    
    logger.info(f"Running command: {' '.join(command)}")
    print(f"Reached at subprocess.run() The output path is {output_dir}")
    
    try:
        HF_CACHE_PATH = "/tmp/huggingface_cache"
        os.makedirs(HF_CACHE_PATH, exist_ok=True)
        env = os.environ.copy()
        env["HF_HOME"] = "/tmp/huggingface"
        env["TRANSFORMERS_CACHE"] = HF_CACHE_PATH
        env["HF_HUB_CACHE"] = HF_CACHE_PATH
        env["NUMBA_CACHE_DIR"] = "/tmp"

        result = subprocess.run(
        command,
        cwd=repo_path,
        check=True,
        capture_output=True,
        text=True,
        env=env)  # pass updated env here
        
        
        logger.info(f"Subprocess STDOUT:\n{result.stdout}")
        logger.info(f"Subprocess STDERR:\n{result.stderr}")

    except subprocess.CalledProcessError as e:
        logger.error(f"Subprocess failed with exit code {e.returncode}")
        logger.error(f"STDOUT:\n{e.stdout}")
        logger.error(f"STDERR:\n{e.stderr}")
        raise Exception(f"Subprocess failed: {e.stderr}")
    
    # Find the generated GLB file
    # The model typically outputs files to a subdirectory named after the input image
    # without the file extension
    base_name = os.path.basename(image_path)
    file_name_without_ext = os.path.splitext(base_name)[0]
    
    # Look for GLB files in the output directory and its subdirectories
    glb_files = glob.glob(os.path.join(output_dir, "**", "*.glb"), recursive=True)
    
    if not glb_files:
        logger.error(f"No GLB files found in {output_dir}")
        raise Exception("No GLB file was generated")
    
    logger.info(f"Found GLB file: {glb_files[0]}")
    print(f"Returned the path {glb_files[0]}")
    return glb_files[0]