import os
import gc
import subprocess
import numpy as np
import torch
from torch.nn import functional as F
from tqdm import tqdm

def is_cuda_usable():
    if not torch.cuda.is_available():
        return False
    try:
        _ = torch.zeros(1, device="cuda")
        return True
    except Exception:
        return False

def clear_vram():
    gc.collect()
    if is_cuda_usable():
        try:
            torch.cuda.empty_cache()
        except Exception:
            pass

def create_classic_boomerang_loop(frames_np):
    """
    Creates classic 100% real-speed full boomerang loop (forward + reverse).
    Reverts to true original speed without artificial slow blending.
    """
    if frames_np is None or len(frames_np) < 4:
        return frames_np

    is_list = isinstance(frames_np, list)
    forward = list(frames_np) if is_list else [f for f in frames_np]
    
    # Exclude boundary duplicates to keep smooth motion flow
    reversed_frames = list(forward[::-1])[1:-1]
    result = forward + reversed_frames
    return result if is_list else np.array(result)

def create_ending_boomerang_loop(frames_np, tail_ratio=0.4, min_tail_frames=24):
    """
    Creates ending-only boomerang loop: plays full video forward normally, 
    then boomerangs only the last 1.5-2.0s tail frames at 100% real speed.
    """
    if frames_np is None or len(frames_np) < 6:
        return frames_np

    is_list = isinstance(frames_np, list)
    forward = list(frames_np) if is_list else [f for f in frames_np]
    n_frames = len(forward)
    
    # Calculate tail frame count (last ~1.5s to 2.0s based on total frames)
    tail_count = max(min_tail_frames, int(n_frames * tail_ratio))
    tail_count = min(n_frames - 2, tail_count)
    
    tail_frames = forward[-tail_count:]
    reversed_tail = list(tail_frames[::-1])[1:-1]
    
    result = forward + reversed_tail
    return result if is_list else np.array(result)

def create_dynamic_boomerang_loop(frames_np, blend_frames=3):
    return create_classic_boomerang_loop(frames_np)

def create_adaptive_speed_ramping(frames_np, multiplier=2):
    """
    Applies non-linear motion-compensated speed ramping (Ease-In / Ease-Out curve).
    Preserves 100% real-time motion speed during fast action/middle segments,
    while smoothly easing start and end keyframes to extend duration naturally.
    """
    if frames_np is None or len(frames_np) < 6:
        return frames_np

    is_list = isinstance(frames_np, list)
    forward = list(frames_np) if is_list else [f for f in frames_np]
    N = len(forward)
    target_count = int((N - 1) * multiplier + 1)

    out_frames = []
    for k in range(target_count):
        s = k / float(target_count - 1)
        # Cubic Smoothstep Easing: 3*s^2 - 2*s^3
        eased_s = s * s * (3.0 - 2.0 * s)
        
        pos = eased_s * (N - 1)
        idx0 = int(pos)
        idx1 = min(N - 1, idx0 + 1)
        alpha = pos - idx0

        if alpha < 0.01 or idx0 == idx1:
            out_frames.append(forward[idx0])
        else:
            f0 = np.array(forward[idx0], dtype=np.float32)
            f1 = np.array(forward[idx1], dtype=np.float32)
            blended = (1.0 - alpha) * f0 + alpha * f1
            out_frames.append(blended.astype(np.uint8) if forward[0].dtype == np.uint8 else blended)

    return out_frames if is_list else np.array(out_frames)

# Download and initialize RIFE Model
if not os.path.exists("RIFEv4.26_0921.zip"):
    print("Downloading RIFE Model...")
    subprocess.run([
        "wget", "-q",
        "https://huggingface.co/thornmaze/RIFE/resolve/main/RIFEv4.26_0921.zip",
        "-O", "RIFEv4.26_0921.zip"
    ], check=True)
    subprocess.run(["unzip", "-o", "RIFEv4.26_0921.zip"], check=True)

from train_log.RIFE_HDv3 import Model
rife_model = Model()
rife_model.load_model("train_log", -1)
rife_model.eval()

@torch.no_grad()
def interpolate_bits(frames_np, multiplier=2, scale=1.0):
    if isinstance(frames_np, list):
        T = len(frames_np)
        H, W, C = frames_np[0].shape
    else:
        T, H, W, C = frames_np.shape

    if multiplier < 2:
        if isinstance(frames_np, np.ndarray):
            return list(frames_np)
        return frames_np

    n_interp = multiplier - 1
    tmp = max(128, int(128 / scale))
    ph = ((H - 1) // tmp + 1) * tmp
    pw = ((W - 1) // tmp + 1) * tmp
    padding = (0, pw - W, 0, ph - H)

    use_cuda = is_cuda_usable()
    curr_device = torch.device("cuda" if use_cuda else "cpu")

    try:
        if hasattr(rife_model, "flownet") and rife_model.flownet is not None:
            rife_model.flownet = rife_model.flownet.to(curr_device)
            if use_cuda:
                rife_model.flownet = rife_model.flownet.half()
            else:
                rife_model.flownet = rife_model.flownet.float()
    except Exception as e:
        print(f"RIFE model device placement notice: {e}")

    def to_tensor(frame_np):
        t = torch.from_numpy(frame_np).to(curr_device)
        t = t.permute(2, 0, 1).unsqueeze(0)
        if curr_device.type == "cuda":
            return F.pad(t, padding).half()
        return F.pad(t, padding).float()

    def from_tensor(tensor):
        t = tensor[0, :, :H, :W]
        t = t.permute(1, 2, 0)
        return t.float().cpu().numpy()

    def make_inference(I0, I1, n):
        if rife_model.version >= 3.9:
            res = []
            for i in range(n):
                res.append(rife_model.inference(I0, I1, (i+1) * 1. / (n+1), scale))
            return res
        else:
            middle = rife_model.inference(I0, I1, scale)
            if n == 1:
                return [middle]
            first_half = make_inference(I0, middle, n=n//2)
            second_half = make_inference(middle, I1, n=n//2)
            if n % 2:
                return [*first_half, middle, *second_half]
            else:
                return [*first_half, *second_half]

    output_frames = []
    I1 = to_tensor(frames_np[0])
    total_steps = T - 1

    with tqdm(total=total_steps, desc="Interpolating", unit="frame") as pbar:
        for i in range(total_steps):
            I0 = I1
            output_frames.append(from_tensor(I0))
            I1 = to_tensor(frames_np[i+1])
            mid_tensors = make_inference(I0, I1, n_interp)
            for mid in mid_tensors:
                output_frames.append(from_tensor(mid))
            if (i + 1) % 50 == 0:
                pbar.update(50)
        pbar.update(total_steps % 50)
        output_frames.append(from_tensor(I1))
    
    del I0, I1, mid_tensors
    if curr_device.type == "cuda" and is_cuda_usable():
        try:
            torch.cuda.empty_cache()
        except Exception:
            pass
    return output_frames


def call_sulphur_rife_api(video_path, multiplier=2, slow_motion=False, upscale=True, enhance_face=False, server_url=None, progress=None):
    """
    Calls Sulphur AI VIP RIFE API (/api/v1/rife-extend) with support for RIFE frame interpolation,
    Real-ESRGAN 1080p HD Super-Resolution, and Temporal Face Sharpening.
    Uses standard library urllib.request for ultra-fast, deadlock-free HTTP transfers.
    Returns path to output MP4 video file.
    """
    import os
    import json
    import tempfile
    import urllib.request
    import config

    if not video_path or not os.path.exists(video_path):
        return None

    target_url = server_url or config.SULPHUR_API_URL or os.environ.get("SULPHUR_API_URL", "http://localhost:6666")
    if not target_url or not str(target_url).strip():
        return None

    clean_url = str(target_url).strip().rstrip("/")
    endpoint = f"{clean_url}/api/v1/rife-extend"

    try:
        print(f"🌐 Calling Sulphur AI VIP RIFE API at {endpoint} (multiplier={multiplier}, slow_motion={slow_motion}, upscale={upscale}, enhance_face={enhance_face})...")
        if progress:
            progress(0.91, desc="🌐 Uploading video to VIP Remote GPU Server...")

        with open(video_path, "rb") as vf:
            file_bytes = vf.read()

        boundary = "----SulphurAIRIFEBoundary7MA4YWxkTrZu0gW"
        body = []
        body.append(f"--{boundary}\r\nContent-Disposition: form-data; name=\"video\"; filename=\"{os.path.basename(video_path)}\"\r\nContent-Type: video/mp4\r\n\r\n".encode("utf-8"))
        body.append(file_bytes)
        body.append(b"\r\n")

        params = [
            ("multiplier", str(int(multiplier))),
            ("slow_motion", "true" if slow_motion else "false"),
            ("upscale", "true" if upscale else "false"),
            ("enhance_face", "true" if enhance_face else "false")
        ]
        for k, v in params:
            body.append(f"--{boundary}\r\nContent-Disposition: form-data; name=\"{k}\"\r\n\r\n{v}\r\n".encode("utf-8"))
        body.append(f"--{boundary}--\r\n".encode("utf-8"))

        payload = b"".join(body)

        req = urllib.request.Request(
            endpoint,
            data=payload,
            headers={
                "Content-Type": f"multipart/form-data; boundary={boundary}",
                "User-Agent": "Mozilla/5.0 (SulphurAI-Space-Client)"
            }
        )

        if progress:
            progress(0.93, desc="💎 Offloading 1080p HD RIFE to RTX 3060 GPU...")

        with urllib.request.urlopen(req, timeout=300) as response:
            res_bytes = response.read()
            ct = response.headers.get("Content-Type", "")

            if "application/json" in ct:
                try:
                    json_data = json.loads(res_bytes.decode("utf-8"))
                    video_url = json_data.get("video_url") or json_data.get("url") or json_data.get("video")
                    if video_url:
                        print(f"✅ VIP RIFE returned direct Video URL: {video_url}")
                        if progress:
                            progress(1.0, desc="✅ VIP 1080p HD Acceleration Complete!")
                        return video_url
                except Exception as e:
                    print(f"Notice parsing JSON video_url: {e}")

            if response.status == 200 and res_bytes:
                if progress:
                    progress(0.97, desc="📥 Writing 1080p HD video result...")
                out_filename = f"vip_rife_{multiplier}x_{os.path.basename(video_path)}"
                out_path = os.path.join(tempfile.gettempdir(), out_filename)
                with open(out_path, "wb") as f:
                    f.write(res_bytes)
                print(f"✅ Sulphur AI VIP RIFE Acceleration succeeded: {out_path}")
                if progress:
                    progress(1.0, desc="✅ VIP 1080p HD Acceleration Complete!")
                return out_path
    except Exception as e:
        print(f"⚠️ VIP RIFE API notice: {e}")

    return None