from fastapi import FastAPI, UploadFile, File, Response, Form, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from PIL import Image
import io
import shutil
from rembg import remove as rembg_remove, new_session
import time
import numpy as np
import tempfile
import uuid
import os
import subprocess
from transformers import pipeline
from transparent_background import Remover
import logging
import asyncio
from datetime import datetime, timedelta
import torch
from .ormbg import ORMBGProcessor
from typing import Dict
from contextlib import contextmanager

from carvekit.ml.files.models_loc import download_all

download_all()

from carvekit.ml.wrap.u2net import U2NET
from carvekit.ml.wrap.basnet import BASNET
from carvekit.ml.wrap.fba_matting import FBAMatting
from carvekit.ml.wrap.deeplab_v3 import DeepLabV3
from carvekit.ml.wrap.tracer_b7 import TracerUniversalB7
from carvekit.api.interface import Interface
from carvekit.pipelines.postprocessing import MattingMethod
from carvekit.pipelines.preprocessing import PreprocessingStub
from carvekit.trimap.generator import TrimapGenerator


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

# Add ORMBG model initialization
ormbg_model_path = os.path.expanduser("~/.ormbg/ormbg.pth")
try:
    ormbg_processor = ORMBGProcessor(ormbg_model_path)
    if torch.cuda.is_available():
        ormbg_processor.to("cuda")
    else:
        ormbg_processor.to("cpu")
except FileNotFoundError:
    logger.error(f"ORMBG model file not found: {ormbg_model_path}")
    print("Error: ORMBG model file not found. Please run 'npm run setup-server' to download it.")
    exit(1)

app = FastAPI()

# Create temp_videos folder if it doesn't exist
TEMP_VIDEOS_DIR = "temp_videos"
os.makedirs(TEMP_VIDEOS_DIR, exist_ok=True)

# Create a frames directory within temp_videos
FRAMES_DIR = os.path.join(TEMP_VIDEOS_DIR, "frames")
os.makedirs(FRAMES_DIR, exist_ok=True)

# Add a dictionary to store processing status
processing_status = {}

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

async def cleanup_old_videos():
    while True:
        current_time = datetime.now()
        for item in os.listdir(TEMP_VIDEOS_DIR):
            item_path = os.path.join(TEMP_VIDEOS_DIR, item)
            item_modified = datetime.fromtimestamp(os.path.getmtime(item_path))
            if current_time - item_modified > timedelta(minutes=10):
                if os.path.isfile(item_path):
                    os.remove(item_path)
                    logger.info(f"Removed old file: {item_path}")
                elif os.path.isdir(item_path):
                    shutil.rmtree(item_path)
                    logger.info(f"Removed old directory: {item_path}")
        await asyncio.sleep(600)  # Run every 10 minutes

# Pre-load all models
bria_model = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device="cpu")
inspyrenet_model = Remover()
inspyrenet_model.model.cpu()
rembg_models = {
    'u2net': new_session('u2net'),
    'u2net_human_seg': new_session('u2net_human_seg'),
    'isnet-general-use': new_session('isnet-general-use'),
    'isnet-anime': new_session('isnet-anime')
}

# Initialize Carvekit models
def initialize_carvekit_model(seg_pipe_class, device='cuda'):
    model = Interface(
        pre_pipe=PreprocessingStub(),
        post_pipe=MattingMethod(
            matting_module=FBAMatting(device=device, input_tensor_size=2048, batch_size=1),
            trimap_generator=TrimapGenerator(),
            device=device
        ),
        seg_pipe=seg_pipe_class(device=device, batch_size=1)
    )
    model.segmentation_pipeline.to('cpu')
    return model

carvekit_models = {
    'u2net': initialize_carvekit_model(U2NET),
    'tracer': initialize_carvekit_model(TracerUniversalB7),
    'basnet': initialize_carvekit_model(BASNET),
    'deeplab': initialize_carvekit_model(DeepLabV3)
}

# ---- Method registry + aliases (NEW) -----------------------------------------
CARVEKIT = list(carvekit_models.keys())          # ['u2net', 'tracer', 'basnet', 'deeplab']
REMBG    = ['u2net_human_seg', 'isnet-general-use', 'isnet-anime']
EXTRA    = ['ormbg', 'bria', 'inspyrenet']
ALLOWED_METHODS = sorted(set(CARVEKIT + REMBG + EXTRA))

ALIASES = {
    'deeplabv3': 'deeplab',
    'tracer-b7': 'tracer',
    'tracer_b7': 'tracer',
    'u2net_human': 'u2net_human_seg',
    'isnet_general_use': 'isnet-general-use',
    'bria14': 'bria',
    # (intentionally no alias for 'rembg'/'open-rmbg')
}

@app.get("/methods")
def list_methods():
    return {"methods": ALLOWED_METHODS}

@app.get("/health")
def health():
    return JSONResponse({"ok": True})
# -----------------------------------------------------------------------------

# Ensure GPU memory is cleared after initialization
torch.cuda.empty_cache()

def process_with_bria(image):
    result = bria_model(image, return_mask=True)
    mask = result
    if not isinstance(mask, Image.Image):
        mask = Image.fromarray((mask * 255).astype('uint8'))
    no_bg_image = Image.new("RGBA", image.size, (0, 0, 0, 0))
    no_bg_image.paste(image, mask=mask)
    return no_bg_image

def process_with_ormbg(image):
    result = ormbg_processor.process_image(image)
    return result

def process_with_inspyrenet(image):
    return inspyrenet_model.process(image, type='rgba')

def process_with_rembg(image, model='u2net'):
    return rembg_remove(image, session=rembg_models[model])

def process_with_carvekit(image, model='u2net'):
    # Initialize segmentation network based on model input
    if model == 'u2net':
        seg_net = U2NET(device='cuda', batch_size=1)
    elif model == 'tracer':
        seg_net = TracerUniversalB7(device='cuda', batch_size=1)
    elif model == 'basnet':
        seg_net = BASNET(device='cuda', batch_size=1)
    elif model == 'deeplab':
        seg_net = DeepLabV3(device='cuda', batch_size=1)
    else:
        raise ValueError("Unsupported model type")

    # Setup the post-processing components
    fba = FBAMatting(device='cuda', input_tensor_size=2048, batch_size=1)
    trimap = TrimapGenerator()
    preprocessing = PreprocessingStub()
    postprocessing = MattingMethod(matting_module=fba, trimap_generator=trimap, device='cuda')

    interface = Interface(pre_pipe=preprocessing, post_pipe=postprocessing, seg_pipe=seg_net)
    processed_image = interface([image])[0]
    
    return processed_image

@contextmanager
def inspyrenet_video_model_context():
    try:
        model = Remover()
        model.model.cuda()
        yield model
    finally:
        model.model.cpu()
        del model
        torch.cuda.empty_cache()

@contextmanager
def carvekit_video_model_context(model_name):
    try:
        if model_name == 'u2net':
            seg_net = U2NET(device='cuda', batch_size=1)
        elif model_name == 'tracer':
            seg_net = TracerUniversalB7(device='cuda', batch_size=1)
        elif model_name == 'basnet':
            seg_net = BASNET(device='cuda', batch_size=1)
        elif model_name == 'deeplab':
            seg_net = DeepLabV3(device='cuda', batch_size=1)
        else:
            raise ValueError("Unsupported model type")

        fba = FBAMatting(device='cuda', input_tensor_size=2048, batch_size=1)
        trimap = TrimapGenerator()
        preprocessing = PreprocessingStub()
        postprocessing = MattingMethod(matting_module=fba, trimap_generator=trimap, device='cuda')

        interface = Interface(pre_pipe=preprocessing, post_pipe=postprocessing, seg_pipe=seg_net)
        yield interface
    finally:
        del seg_net, fba, trimap, preprocessing, postprocessing, interface
        torch.cuda.empty_cache()


# Create a global lock for GPU operations
gpu_lock = asyncio.Lock()

@app.post("/remove_background/")
async def remove_background(file: UploadFile = File(...), method: str = Form(...)):
    try:
        # normalize aliases
        method = ALIASES.get(method, method)

        # early reject unknown methods with 400, not 500
        if method not in ALLOWED_METHODS:
            raise HTTPException(status_code=400, detail="Invalid method")

        image_data = await file.read()
        image = Image.open(io.BytesIO(image_data)).convert('RGB')
        
        start_time = time.time()

        async def process_image():
            if method == 'bria':
                return await asyncio.to_thread(process_with_bria, image)
            elif method == 'inspyrenet':
                async with gpu_lock:
                    try:
                        inspyrenet_model.model.to('cuda')
                        result = await asyncio.to_thread(inspyrenet_model.process, image, type='rgba')
                    finally:
                        inspyrenet_model.model.to('cpu')
                    return result
            elif method in ['u2net_human_seg', 'isnet-general-use', 'isnet-anime']:
                return await asyncio.to_thread(process_with_rembg, image, model=method)
            elif method == 'ormbg':
                return await asyncio.to_thread(process_with_ormbg, image)
            elif method in ['u2net', 'tracer', 'basnet', 'deeplab']:
                async with gpu_lock:
                    try:
                        carvekit_models[method].segmentation_pipeline.to('cuda')
                        result = await asyncio.to_thread(carvekit_models[method], [image])
                    finally:
                        carvekit_models[method].segmentation_pipeline.to('cpu')
                    return result[0]
            else:
                raise HTTPException(status_code=400, detail="Invalid method")

        no_bg_image = await process_image()
        
        process_time = time.time() - start_time
        print(f"Background removal time ({method}): {process_time:.2f} seconds")
        
        async with gpu_lock:
            torch.cuda.empty_cache()
        
        with io.BytesIO() as output:
            no_bg_image.save(output, format="PNG")
            content = output.getvalue()

        return Response(content=content, media_type="image/png")

    except HTTPException as he:
        # let 4xx bubble as-is
        raise he
    except Exception as e:
        print(str(e))
        raise HTTPException(status_code=500, detail=str(e))

async def process_frame(frame_path, method):
    img = Image.open(frame_path).convert('RGB')
    
    if method == 'bria':
        processed_frame = await asyncio.to_thread(process_with_bria, img)
    elif method in ['u2net_human_seg', 'isnet-general-use', 'isnet-anime']:
        processed_frame = await asyncio.to_thread(process_with_rembg, img, model=method)
    elif method == 'ormbg':
        processed_frame = await asyncio.to_thread(process_with_ormbg, img)
    else:
        raise ValueError("Invalid method")
    
    return processed_frame

async def process_video(video_path, method, video_id):
    try:
        processing_status[video_id] = {'status': 'processing', 'progress': 0, 'message': 'Initializing'}
        
        logger.info(f"Starting video processing: {video_path}")
        logger.info(f"Method: {method}")
        logger.info(f"Video ID: {video_id}")


        # Check video frame count
        frame_count_command = ['ffmpeg.ffprobe', '-v', 'error', '-select_streams', 'v:0', '-count_packets', 
                               '-show_entries', 'stream=nb_read_packets', '-of', 'csv=p=0', video_path]
        process = await asyncio.create_subprocess_exec(
            *frame_count_command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await process.communicate()
        
        if process.returncode != 0:
            logger.error(f"Error counting frames: {stderr.decode()}")
            processing_status[video_id] = {'status': 'error', 'message': 'Error counting frames'}
            return

        frame_count = int(stdout.decode().strip())
        logger.info(f"Video frame count: {frame_count}")

        #DISABLED VIDEO LENGTH LIMIT
        #if frame_count > 250:
        #    logger.warning(f"Video too long: {frame_count} frames")
        #    processing_status[video_id] = {'status': 'error', 'message': 'Video too long (max 250 frames)'}
        #    return

        # Create a unique directory for this video's frames
        frames_dir = os.path.join(FRAMES_DIR, video_id)
        os.makedirs(frames_dir, exist_ok=True)
        logger.info(f"Created frames directory: {frames_dir}")

        # Extract frames from video
        processing_status[video_id] = {'status': 'processing', 'progress': 0, 'message': 'Extracting frames'}
        extract_command = ['ffmpeg', '-i', video_path, f'{frames_dir}/frame_%05d.png']
        logger.info(f"Executing frame extraction command: {' '.join(extract_command)}")
        process = await asyncio.create_subprocess_exec(
            *extract_command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await process.communicate()
        
        if process.returncode != 0:
            logger.error(f"Error extracting frames: {stderr.decode()}")
            processing_status[video_id] = {'status': 'error', 'message': 'Error extracting frames'}
            return

        # Process frames
        processing_status[video_id] = {'status': 'processing', 'progress': 0, 'message': 'Removing background'}
        frame_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.png')])
        total_frames = len(frame_files)
        logger.info(f"Number of extracted frames: {total_frames}")

        if total_frames == 0:
            logger.error("No frames were extracted from the video")
            processing_status[video_id] = {'status': 'error', 'message': 'No frames were extracted from the video'}
            return

        # Initialize the model once, outside the batch processing loop
        if method == 'inspyrenet':
            print("start init")
            model_context = inspyrenet_video_model_context()
            print("start enter")
            model = model_context.__enter__()
            print("finish enter")
        elif method in ['u2net', 'tracer', 'basnet', 'deeplab']:
            model_context = carvekit_video_model_context(method)
            model = model_context.__enter__()
        else:
            model = None  # For other methods that don't require a specific model

        try:
            async def process_frame_batch(start_idx, end_idx):
                for i in range(start_idx, min(end_idx, total_frames)):
                    frame_file = frame_files[i]
                    frame_path = os.path.join(frames_dir, frame_file)
                    img = Image.open(frame_path).convert('RGB')

                    if method == 'inspyrenet':
                        processed_frame = model.process(img, type='rgba')
                    elif method in ['u2net', 'tracer', 'basnet', 'deeplab']:
                        processed_frame = model([img])[0]
                    else:
                        processed_frame = await process_frame(frame_path, method)

                    processed_frame.save(frame_path, format='PNG')
                    progress = (i + 1) / total_frames * 100
                    processing_status[video_id] = {'status': 'processing', 'progress': progress}

            batch_size = 3
            for i in range(0, total_frames, batch_size):
                await process_frame_batch(i, i + batch_size)
                await asyncio.sleep(0)  # Allow other tasks to run

        finally:
            # Ensure we clean up the model context
            if method in ['inspyrenet', 'u2net', 'tracer', 'basnet', 'deeplab']:
                model_context.__exit__(None, None, None)

        # Create output video
        processing_status[video_id] = {'status': 'processing', 'progress': 100, 'message': 'Encoding video'}
        output_path = os.path.join(TEMP_VIDEOS_DIR, f"output_{video_id}.webm")
        create_video_command = [
            'ffmpeg',
            '-framerate', '24',
            '-i', f'{frames_dir}/frame_%05d.png',
            '-c:v', 'libvpx-vp9',
            '-pix_fmt', 'yuva420p',
            '-lossless', '1',
            output_path
        ]
        logger.info(f"Executing video creation command: {' '.join(create_video_command)}")
        process = await asyncio.create_subprocess_exec(
            *create_video_command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await process.communicate()
        
        if process.returncode != 0:
            logger.error(f"Error creating output video: {stderr.decode()}")
            processing_status[video_id] = {'status': 'error', 'message': 'Error creating output video'}
            return

        logger.info(f"Video processing completed. Output path: {output_path}")
        processing_status[video_id] = {'status': 'completed', 'output_path': output_path}

    except Exception as e:
        logger.exception("Error in video processing")
        processing_status[video_id] = {'status': 'error', 'message': str(e)}
    finally:
        torch.cuda.empty_cache()

        # Clean up frames directory
        for file in os.listdir(frames_dir):
            os.remove(os.path.join(frames_dir, file))
        os.rmdir(frames_dir)
        logger.info(f"Cleaned up frames directory: {frames_dir}")

@app.post("/remove_background_video/")
async def remove_background_video(background_tasks: BackgroundTasks, file: UploadFile = File(...), method: str = Form(...)):
    try:
        logger.info(f"Starting video background removal with method: {method}")
        
        # Generate a unique filename for the uploaded video
        video_id = str(uuid.uuid4())
        filename = f"input_{video_id}.mp4"
        file_path = os.path.join(TEMP_VIDEOS_DIR, filename)
        
        # Save uploaded video to the temp_videos folder
        with open(file_path, "wb") as buffer:
            content = await file.read()
            buffer.write(content)

        logger.info(f"Video file saved: {file_path}")
        logger.info(f"File exists: {os.path.exists(file_path)}")
        logger.info(f"File size: {os.path.getsize(file_path)} bytes")

        if not os.path.exists(file_path):
            raise HTTPException(status_code=500, detail=f"Failed to create video file: {file_path}")

        # Start processing in the background
        background_tasks.add_task(process_video, file_path, method, video_id)
        
        return {"video_id": video_id}

    except Exception as e:
        logger.exception(f"Error in video processing: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error in video processing: {str(e)}")

@app.get("/status/{video_id}")
async def get_status(video_id: str):
    if video_id not in processing_status:
        raise HTTPException(status_code=404, detail="Video ID not found")
    
    status = processing_status[video_id]
    
    if status['status'] == 'completed':
        output_path = status['output_path']
        if not os.path.exists(output_path):
            raise HTTPException(status_code=404, detail="Processed video file not found")
        
        return FileResponse(output_path, media_type="video/webm", filename=f"processed_video_{video_id}.webm")
    
    return status

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(cleanup_old_videos())
    

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=9876)