import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from PIL import Image 
import numpy as np 
import os 
import tempfile
import spaces
import gradio as gr
import subprocess
import sys
import cv2
import threading
import queue
import time
from collections import deque
from deep_translator import GoogleTranslator

def install_flash_attn_wheel():
    flash_attn_wheel_url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+cu123torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", flash_attn_wheel_url])
        print("Wheel installed successfully!")
    except subprocess.CalledProcessError as e:
        print(f"Failed to install the flash attnetion wheel. Error: {e}")

install_flash_attn_wheel()

try:
    from mmengine.visualization import Visualizer
except ImportError:
    Visualizer = None
    print("Warning: mmengine is not installed, visualization is disabled.")
    
# Load the model and tokenizer 
model_path = "ByteDance/Sa2VA-4B"
 
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype="auto",
    device_map="cuda:0",
    trust_remote_code=True,
).eval().cuda()

tokenizer = AutoTokenizer.from_pretrained(
    model_path,
    trust_remote_code = True,
)

class WebcamProcessor:
    def __init__(self, model, tokenizer, fps_target=15, buffer_size=5):
        self.model = model
        self.tokenizer = tokenizer
        self.fps_target = fps_target
        self.frame_interval = 1.0 / fps_target
        self.buffer_size = buffer_size
        self.frame_buffer = deque(maxlen=buffer_size)
        self.result_queue = queue.Queue()
        self.is_running = False
        self.last_process_time = 0
        
    def start(self):
        self.is_running = True
        self.capture = cv2.VideoCapture(0)
        self.capture_thread = threading.Thread(target=self._capture_loop)
        self.process_thread = threading.Thread(target=self._process_loop)
        self.capture_thread.start()
        self.process_thread.start()
        
    def stop(self):
        self.is_running = False
        if hasattr(self, 'capture_thread'):
            self.capture_thread.join()
            self.process_thread.join()
            self.capture.release()
            
    def _capture_loop(self):
        while self.is_running:
            ret, frame = self.capture.read()
            if ret:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frame = cv2.resize(frame, (640, 480))
                current_time = time.time()
                if current_time - self.last_process_time >= self.frame_interval:
                    self.frame_buffer.append(frame)
                    self.last_process_time = current_time
                    
    def _process_loop(self):
        while self.is_running:
            if len(self.frame_buffer) >= self.buffer_size:
                frames = list(self.frame_buffer)
                try:
                    result = self.model.predict_forward(
                        video=frames,
                        text="<image>Describe what you see",
                        tokenizer=self.tokenizer
                    )
                    self.result_queue.put(result)
                except Exception as e:
                    print(f"Processing error: {e}")
                self.frame_buffer.clear()
            time.sleep(0.1)

from third_parts import VideoReader
def read_video(video_path, video_interval):
    vid_frames = VideoReader(video_path)[::video_interval]
    
    temp_dir = tempfile.mkdtemp()
    os.makedirs(temp_dir, exist_ok=True)
    image_paths = []
    
    for frame_idx in range(len(vid_frames)):
        frame_image = vid_frames[frame_idx]
        frame_image = frame_image[..., ::-1]
        frame_image = Image.fromarray(frame_image)
        vid_frames[frame_idx] = frame_image
        image_path = os.path.join(temp_dir, f"frame_{frame_idx:04d}.jpg")
        frame_image.save(image_path, format="JPEG")
        image_paths.append(image_path)
    return vid_frames, image_paths

def visualize(pred_mask, image_path, work_dir):
    visualizer = Visualizer()
    img = cv2.imread(image_path)
    visualizer.set_image(img)
    visualizer.draw_binary_masks(pred_mask, colors='g', alphas=0.4)
    visual_result = visualizer.get_image()
    output_path = os.path.join(work_dir, os.path.basename(image_path))
    cv2.imwrite(output_path, visual_result)
    return output_path

def translate_to_korean(text):
    try:
        translator = GoogleTranslator(source='en', target='ko')
        return translator.translate(text)
    except Exception as e:
        print(f"Translation error: {e}")
        return text

@spaces.GPU
def image_vision(image_input_path, prompt):
    is_korean = any(ord('가') <= ord(char) <= ord('힣') for char in prompt)
    
    image_path = image_input_path
    text_prompts = f"<image>{prompt}"
    image = Image.open(image_path).convert('RGB')
    input_dict = {
        'image': image,
        'text': text_prompts,
        'past_text': '',
        'mask_prompts': None,
        'tokenizer': tokenizer,
    }
    return_dict = model.predict_forward(**input_dict)
    print(return_dict)
    answer = return_dict["prediction"]
    
    if is_korean:
        if '[SEG]' in answer:
            parts = answer.split('[SEG]')
            translated_parts = [translate_to_korean(part.strip()) for part in parts]
            answer = '[SEG]'.join(translated_parts)
        else:
            answer = translate_to_korean(answer)
    
    seg_image = return_dict["prediction_masks"]
    
    if '[SEG]' in answer and Visualizer is not None:
        pred_masks = seg_image[0]
        temp_dir = tempfile.mkdtemp()
        pred_mask = pred_masks
        os.makedirs(temp_dir, exist_ok=True)
        seg_result = visualize(pred_mask, image_input_path, temp_dir)
        return answer, seg_result
    else:
        return answer, None

@spaces.GPU(duration=80)
def video_vision(video_input_path, prompt, video_interval):
    is_korean = any(ord('가') <= ord(char) <= ord('힣') for char in prompt)
    
    cap = cv2.VideoCapture(video_input_path)
    original_fps = cap.get(cv2.CAP_PROP_FPS)
    frame_skip_factor = video_interval
    new_fps = original_fps / frame_skip_factor

    vid_frames, image_paths = read_video(video_input_path, video_interval)
    question = f"<image>{prompt}"
    result = model.predict_forward(
        video=vid_frames,
        text=question,
        tokenizer=tokenizer,
    )
    prediction = result['prediction']
    print(prediction)

    if is_korean:
        if '[SEG]' in prediction:
            parts = prediction.split('[SEG]')
            translated_parts = [translate_to_korean(part.strip()) for part in parts]
            prediction = '[SEG]'.join(translated_parts)
        else:
            prediction = translate_to_korean(prediction)

    if '[SEG]' in prediction and Visualizer is not None:
        _seg_idx = 0
        pred_masks = result['prediction_masks'][_seg_idx]
        seg_frames = []
        for frame_idx in range(len(vid_frames)):
            pred_mask = pred_masks[frame_idx]
            temp_dir = tempfile.mkdtemp()
            os.makedirs(temp_dir, exist_ok=True)
            seg_frame = visualize(pred_mask, image_paths[frame_idx], temp_dir)
            seg_frames.append(seg_frame)

        output_video = "output_video.mp4"
        frame = cv2.imread(seg_frames[0])
        height, width, layers = frame.shape
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video = cv2.VideoWriter(output_video, fourcc, new_fps, (width, height))

        for img_path in seg_frames:
            frame = cv2.imread(img_path)
            video.write(frame)

        video.release()
        print(f"Video created successfully at {output_video}")

        return prediction, output_video
    else:
        return prediction, None

@spaces.GPU
def webcam_vision(prompt):
    is_korean = any(ord('가') <= ord(char) <= ord('힣') for char in prompt)
    
    if not hasattr(webcam_vision, 'processor'):
        webcam_vision.processor = WebcamProcessor(model, tokenizer)
    
    if not webcam_vision.processor.is_running:
        webcam_vision.processor.start()
    
    try:
        result = webcam_vision.processor.result_queue.get(timeout=5)
        prediction = result['prediction']
        
        if is_korean:
            prediction = translate_to_korean(prediction)
            
        return prediction
    except queue.Empty:
        return "No results available yet"
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
with gr.Blocks(analytics_enabled=False) as demo:
    with gr.Column():
        gr.Markdown("# Sa2VA: Marrying SAM2 with LLaVA for Dense Grounded Understanding of Images and Videos")

        
        with gr.Tab("Single Image"):
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(label="Image IN", type="filepath")
                    with gr.Row():
                        instruction = gr.Textbox(label="Instruction", scale=4)
                        submit_image_btn = gr.Button("Submit", scale=1)
                with gr.Column():
                    output_res = gr.Textbox(label="Response")
                    output_image = gr.Image(label="Segmentation", type="numpy")
    
            submit_image_btn.click(
                fn = image_vision,
                inputs = [image_input, instruction],
                outputs = [output_res, output_image]
            )
            
        with gr.Tab("Video"):
            with gr.Row():
                with gr.Column():
                    video_input = gr.Video(label="Video IN")
                    frame_interval = gr.Slider(label="Frame interval", step=1, minimum=1, maximum=12, value=6)
                    with gr.Row():
                        vid_instruction = gr.Textbox(label="Instruction", scale=4)
                        submit_video_btn = gr.Button("Submit", scale=1)
                with gr.Column():
                    vid_output_res = gr.Textbox(label="Response")
                    output_video = gr.Video(label="Segmentation")
            
            submit_video_btn.click(
                fn = video_vision,
                inputs = [video_input, vid_instruction, frame_interval],
                outputs = [vid_output_res, output_video]
            )
            
        with gr.Tab("Webcam"):
            with gr.Row():
                with gr.Column():
                    webcam_input = gr.Image(source="webcam", streaming=True)
                    with gr.Row():
                        webcam_instruction = gr.Textbox(
                            label="Instruction",
                            placeholder="Enter instruction here...",
                            scale=4
                        )
                        start_button = gr.Button("Start", scale=1)
                        stop_button = gr.Button("Stop", scale=1)
                with gr.Column():
                    webcam_output = gr.Textbox(label="Response")
                    processed_view = gr.Image(label="Processed View")
            
            status_text = gr.Textbox(label="Status", value="Ready")
            
            start_button.click(
                fn=lambda x: webcam_vision(x),
                inputs=[webcam_instruction],
                outputs=[webcam_output]
            )
            
            stop_button.click(
                fn=lambda: "Stopped" if hasattr(webcam_vision, 'processor') and webcam_vision.processor.stop() else "Not running",
                outputs=[status_text]
            )

demo.queue().launch(show_api=False, show_error=True)