import os
import subprocess
import torch
import torch.nn as nn
import cv2
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks
from fastapi.responses import FileResponse, JSONResponse
import shutil
import uuid
import urllib.request
from concurrent.futures import ThreadPoolExecutor

app = FastAPI(title="Ultra Fast AI Video Enhancer (Lanczos4 + RIFE-PyTorch Parallel)")

UPLOAD_DIR = "/tmp/uploads"
OUTPUT_DIR = "/tmp/outputs"
MODEL_DIR = "/app/weights"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(MODEL_DIR, exist_ok=True)

RIFE_URL = "https://huggingface.co/hfmaster/models-moved/resolve/main/rife/rife49.pth"
RIFE_PATH = os.path.join(MODEL_DIR, "rife49.pth")

device = torch.device('cpu')
rife_model = None

tasks_db = {}
# Thread pool untuk paralelisme CPU level tinggi
executor = ThreadPoolExecutor(max_workers=16)

# ==============================================================================
# 1. PURE PYTORCH RIFE MODEL DEFINITION
# ==============================================================================
def warp(tenInput, tenFlow):
    backwarp_tenGrid = {}
    k = str(tenFlow.device) + '_' + str(tenFlow.size())
    if k not in backwarp_tenGrid:
        gX, gY = torch.meshgrid(
            torch.arange(0, tenFlow.size(3), device=tenFlow.device),
            torch.arange(0, tenFlow.size(2), device=tenFlow.device),
            indexing='xy'
        )
        backwarp_tenGrid[k] = torch.stack((gX, gY), 2).float()
    
    tenGrid = backwarp_tenGrid[k]
    tenFlow = torch.cat([
        tenFlow[:, 0:1, :, :] / ((tenInput.size(3) - 1.0) / 2.0),
        tenFlow[:, 1:2, :, :] / ((tenInput.size(2) - 1.0) / 2.0)
    ], 1)
    
    g = (tenGrid + tenFlow.permute(0, 2, 3, 1)) / (torch.tensor([[[[tenInput.size(3) - 1.0, tenInput.size(2) - 1.0]]]], device=tenFlow.device) / 2.0) - 1.0
    return nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True)

class RIFEFlowNet(nn.Module):
    def __init__(self):
        super(RIFEFlowNet, self).__init__()
        self.block0 = nn.Sequential(
            nn.Conv2d(6, 32, 3, 1, 1),
            nn.PReLU(32),
            nn.Conv2d(32, 32, 3, 1, 1),
            nn.PReLU(32)
        )
        self.conv_flow = nn.Conv2d(32, 4, 3, 1, 1)

    def forward(self, img0, img1):
        x = torch.cat([img0, img1], 1)
        x = self.block0(x)
        flow = self.conv_flow(x)
        return flow[:, :2], flow[:, 2:]

class RIFENet(nn.Module):
    def __init__(self):
        super(RIFENet, self).__init__()
        self.flownet = RIFEFlowNet()
        self.unet = nn.Sequential(
            nn.Conv2d(12, 32, 3, 1, 1),
            nn.PReLU(32),
            nn.Conv2d(32, 3, 3, 1, 1)
        )

    def forward(self, img0, img1, timestep=0.5):
        flow_01, flow_10 = self.flownet(img0, img1)
        warped_img0 = warp(img0, flow_01 * timestep)
        warped_img1 = warp(img1, flow_10 * (1.0 - timestep))
        merged = (1.0 - timestep) * warped_img0 + timestep * warped_img1
        x = torch.cat([img0, img1, warped_img0, warped_img1], 1)
        refinement = self.unet(x)
        return torch.clamp(merged + refinement, 0.0, 1.0)

def init_models(task_id):
    global rife_model
    if rife_model is None:
        tasks_db[task_id]["logs"].append("[DEBUG] Memeriksa keberadaan file bobot RIFE (rife49.pth)...")
        if not os.path.exists(RIFE_PATH):
            tasks_db[task_id]["logs"].append("[DEBUG] File RIFE tidak ditemukan. Mengunduh weights dari Hugging Face...")
            urllib.request.urlretrieve(RIFE_URL, RIFE_PATH)
            tasks_db[task_id]["logs"].append("[DEBUG] Bobot model RIFE sukses diunduh.")
        else:
            tasks_db[task_id]["logs"].append("[DEBUG] Berkas bobot RIFE terdeteksi di cache.")
            
        rife_model = RIFENet()
        state_dict = torch.load(RIFE_PATH, map_location=device)
        cleaned_state = {k.replace("module.", ""): v for k, v in state_dict.items()}
        rife_model.load_state_dict(cleaned_state, strict=False)
        rife_model.eval().to(device)
        tasks_db[task_id]["logs"].append("[DEBUG] Model RIFE PyTorch sukses dimuat ke CPU RAM.")

# ==============================================================================
# 2. FASTAPI GATEWAY ENDPOINTS
# ==============================================================================
@app.get("/")
def read_root():
    return {"status": "online", "message": "Ultra Fast RIFE-PyTorch Engine is running."}

@app.get("/status/{task_id}")
def check_status(task_id: str):
    if task_id not in tasks_db:
        return JSONResponse(status_code=404, content={"error": "Task not found"})
    return tasks_db[task_id]

@app.get("/download/{task_id}")
def download_result(task_id: str):
    if task_id not in tasks_db:
        return JSONResponse(status_code=404, content={"error": "Task not found"})
    task = tasks_db[task_id]
    if task["status"] != "completed":
        return JSONResponse(status_code=400, content={"error": f"Task state is: {task['status']}"})
    return FileResponse(task["output_file"], media_type="video/mp4", filename=f"ai_enhanced_{task_id}.mp4")

# ==============================================================================
# 3. HIGH PERFORMANCE WORKERS
# ==============================================================================
def process_single_frame_fast(frame, out_w, out_h):
    # Lanczos4 upscaling + Saturation boost 1.35x
    resized = cv2.resize(frame, (out_w, out_h), interpolation=cv2.INTER_LANCZOS4)
    hsv = cv2.cvtColor(resized, cv2.COLOR_BGR2HSV)
    hsv[:, :, 1] = cv2.multiply(hsv[:, :, 1], 1.35)
    enhanced = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
    
    img_float = enhanced.astype(np.float32) / 255.0
    return img_float

def interpolate_frame_rife(prev_np, curr_np, t, model_local, dev_local):
    with torch.no_grad():
        t1 = torch.from_numpy(np.transpose(prev_np, (2, 0, 1))).float().unsqueeze(0).to(dev_local)
        t2 = torch.from_numpy(np.transpose(curr_np, (2, 0, 1))).float().unsqueeze(0).to(dev_local)
        inter_tensor = model_local(t1, t2, timestep=t)
        inter_np = inter_tensor.squeeze().float().cpu().clamp_(0, 1).numpy()
        inter_np = np.transpose(inter_np, (1, 2, 0))
        return (inter_np * 255.0).round().astype(np.uint8)

def process_video_worker(task_id: str, input_path: str, output_path: str, crf: int):
    try:
        tasks_db[task_id]["status"] = "processing"
        tasks_db[task_id]["logs"] = ["[DEBUG] Inisialisasi pengolahan video asinkron..."]
        
        init_models(task_id)
        
        cap = cv2.VideoCapture(input_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if fps <= 0:
            fps = 30.0
            
        tasks_db[task_id]["total_frames"] = total_frames
        tasks_db[task_id]["logs"].append(f"[DEBUG] Informasi video: {width}x{height} @ {fps} FPS. Total frame: {total_frames}")
        
        # Target 4K (3840x2160 atau 2160x3840 jika vertikal)
        out_w, out_h = 3840, 2160
        if height > width:
            out_w, out_h = 2160, 3840
            
        target_fps = fps * 4
        tasks_db[task_id]["logs"].append(f"[DEBUG] Target resolusi output: {out_w}x{out_h}. Target framerate: {target_fps} FPS.")
        
        # Pipeline FFmpeg input rawvideo RGB24 via standard input
        ffmpeg_cmd = [
            "ffmpeg", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24", 
            "-s", f"{out_w}x{out_h}", "-r", str(target_fps), "-i", "-",
            "-i", input_path, "-map", "0:v:0", "-map", "1:a:0?",
            "-c:v", "libx264", "-crf", str(crf), "-preset", "ultrafast",
            "-pix_fmt", "yuv420p", output_path
        ]
        
        tasks_db[task_id]["logs"].append("[DEBUG] Menginisialisasi pipa data subproses encoder FFmpeg...")
        pipe = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
        
        frame_idx = 0
        prev_float_np = None
        frames_buffer = []
        
        tasks_db[task_id]["logs"].append("[DEBUG] Memulai pemrosesan paralel spasial upscale (Lanczos4) + temporal interpolation (RIFE)...")
        
        with torch.no_grad():
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                frames_buffer.append(frame)
                
                # Proses dalam kelompok 8 frame untuk meningkatkan concurrency multi-thread
                if len(frames_buffer) >= 8:
                    futures_sr = [executor.submit(process_single_frame_fast, f, out_w, out_h) for f in frames_buffer]
                    enhanced_float_nps = [fut.result() for fut in futures_sr]
                    
                    for curr_float_np in enhanced_float_nps:
                        if prev_float_np is not None:
                            # Tiga frame interpolasi (untuk melipatgandakan 4x lipat FPS) diproses secara paralel
                            futures_rife = [
                                executor.submit(interpolate_frame_rife, prev_float_np, curr_float_np, t, rife_model, device) 
                                for t in [0.25, 0.50, 0.75]
                            ]
                            inter_frames = [fut.result() for fut in futures_rife]
                            for inter_frame in inter_frames:
                                pipe.stdin.write(inter_frame.tobytes())
                        
                        cur_frame = (curr_float_np * 255.0).round().astype(np.uint8)
                        pipe.stdin.write(cur_frame.tobytes())
                        
                        prev_float_np = curr_float_np
                        frame_idx += 1
                        tasks_db[task_id]["processed_frames"] = frame_idx
                        
                        if frame_idx % 20 == 0 or frame_idx == total_frames:
                            tasks_db[task_id]["logs"].append(f"[DEBUG] Progres pengolahan: {frame_idx}/{total_frames} frame asli sukses diselesaikan.")
                    
                    frames_buffer = []
            
            # Sisa sisa frame di buffer buffer
            if len(frames_buffer) > 0:
                for frame in frames_buffer:
                    curr_float_np = process_single_frame_fast(frame, out_w, out_h)
                    if prev_float_np is not None:
                        futures_rife = [
                            executor.submit(interpolate_frame_rife, prev_float_np, curr_float_np, t, rife_model, device) 
                            for t in [0.25, 0.50, 0.75]
                        ]
                        inter_frames = [fut.result() for fut in futures_rife]
                        for inter_frame in inter_frames:
                            pipe.stdin.write(inter_frame.tobytes())
                    
                    cur_frame = (curr_float_np * 255.0).round().astype(np.uint8)
                    pipe.stdin.write(cur_frame.tobytes())
                    prev_float_np = curr_float_np
                    frame_idx += 1
                    tasks_db[task_id]["processed_frames"] = frame_idx
                    tasks_db[task_id]["logs"].append(f"[DEBUG] Progres pengolahan akhir: {frame_idx}/{total_frames} frame selesai.")
                    
        cap.release()
        pipe.stdin.close()
        pipe.wait()
        
        # Bersihkan berkas input temporer
        try:
            if os.path.exists(input_path):
                os.remove(input_path)
        except:
            pass
            
        tasks_db[task_id]["status"] = "completed"
        tasks_db[task_id]["logs"].append("[DEBUG] Pengkodean video sukses diselesaikan! File MP4 output 120 FPS siap diunduh.")
        print(f"[{task_id}] Enhancement Completed successfully.")
        
    except Exception as e:
        tasks_db[task_id]["status"] = "failed"
        tasks_db[task_id]["error"] = str(e)
        if "logs" in tasks_db[task_id]:
            tasks_db[task_id]["logs"].append(f"[FATAL ERROR] Gagal mengolah video: {str(e)}")

@app.post("/enhance")
async def enhance_video(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    crf: int = Form(20)
):
    task_id = str(uuid.uuid4())
    input_path = os.path.join(UPLOAD_DIR, f"{task_id}_{file.filename}")
    with open(input_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    final_output_mp4 = os.path.join(OUTPUT_DIR, f"{task_id}_enhanced.mp4")
    
    tasks_db[task_id] = {
        "status": "queued",
        "processed_frames": 0,
        "total_frames": 0,
        "output_file": final_output_mp4,
        "logs": ["[DEBUG] Menerima berkas video masukan...", f"[DEBUG] Membuat ID tugas baru: {task_id}"]
    }
    background_tasks.add_task(process_video_worker, task_id, input_path, final_output_mp4, crf)
    return {
        "status": "queued",
        "task_id": task_id,
        "endpoints": {"status": f"/status/{task_id}", "download": f"/download/{task_id}"}
    }