import os
import shlex
import sys
from pathlib import Path
import numpy as np
import subprocess
import cv2

import preproc.launch_hamer as hamer
from preproc.launch_slam import split_frames_shots, get_command, check_intrins
from preproc.extract_frames import video_to_frames, split_frame


def is_nonempty(d):
    return os.path.isdir(d) and len(os.listdir(d)) > 0


def preprocess_frames(img_dir, src_path, overwrite=False, **kwargs):
    if not overwrite and is_nonempty(img_dir):
        print(f"FOUND {len(os.listdir(img_dir))} FRAMES in {img_dir}")
        return
    print(f"EXTRACTING FRAMES FROM {src_path} TO {img_dir}")
    print(kwargs)

    # out = video_to_frames(src_path, img_dir, overwrite=overwrite, **kwargs)
    out = split_frame(src_path, img_dir, overwrite=overwrite, **kwargs)
    assert out == 0, "FAILED FRAME EXTRACTION"


def preprocess_tracks(datatype, img_dir, track_dir, shot_dir, gpu, overwrite=False):
    """
    :param img_dir
    :param track_dir, expected format: res_root/track_name/sequence
    :param shot_dir, expected format: res_root/shot_name/sequence
    """
    if not overwrite and is_nonempty(track_dir):
        print(f"FOUND TRACKS IN {track_dir}")
        return

    print(f"RUNNING HAMER ON {img_dir}")
    track_root, seq = os.path.split(track_dir.rstrip("/"))
    res_root, track_name = os.path.split(track_root)
    shot_name = shot_dir.rstrip("/").split("/")[-2]

    hamer.process_seq(
        [gpu],
        res_root,
        seq,
        img_dir,
        track_name=track_name,
        shot_name=shot_name,
        datatype=datatype,
        overwrite=overwrite,
    )


def load_vipe_cameras(vipe_dir, seq_name, img_dir, start=0, end=-1):
    """
    Load VIPE camera outputs and convert to DROID-SLAM format
    
    Args:
        vipe_dir: Path to VIPE results directory
        seq_name: Sequence name
        img_dir: Image directory to get image size
        start: Start frame index
        end: End frame index
        
    Returns:
        w2c: (N, 4, 4) world-to-camera matrices
        intrins_full: (N, 6) intrinsics [fx, fy, cx, cy, W, H]
    """
    # Load VIPE pose and intrinsics
    vipe_pose_path = os.path.join(vipe_dir, "pose", f"{seq_name}.npz")
    vipe_intrins_path = os.path.join(vipe_dir, "intrinsics", f"{seq_name}.npz")
    
    if not os.path.exists(vipe_pose_path):
        raise FileNotFoundError(f"VIPE pose file not found: {vipe_pose_path}")
    if not os.path.exists(vipe_intrins_path):
        raise FileNotFoundError(f"VIPE intrinsics file not found: {vipe_intrins_path}")
    
    print(f"Loading VIPE cameras from {vipe_dir}")
    pose_data = np.load(vipe_pose_path)
    c2w = pose_data['data']  # (N, 4, 4) camera-to-world
    pose_inds = pose_data['inds']  # (N,) frame indices
    
    intrins_data = np.load(vipe_intrins_path)
    intrins = intrins_data['data']  # (N, 4) [fx, fy, cx, cy]
    intrins_inds = intrins_data['inds']  # (N,) frame indices
    
    # Verify indices match
    assert np.array_equal(pose_inds, intrins_inds), "Pose and intrinsics indices don't match!"
    
    # Get image size from first image
    image_files = sorted([f for f in os.listdir(img_dir) 
                         if f.endswith(('.png', '.jpg', '.jpeg'))])
    if not image_files:
        # Infer from intrinsics (cx, cy should be roughly at center)
        img_width = int(intrins[0, 2] * 2)
        img_height = int(intrins[0, 3] * 2)
        print(f"Inferred image size from intrinsics: {img_width}x{img_height}")
    else:
        img_path = os.path.join(img_dir, image_files[0])
        img = cv2.imread(img_path)
        img_height, img_width = img.shape[:2]
        print(f"Got image size from images: {img_width}x{img_height}")
    
    # Select frames based on start/end
    if end < 0:
        end = len(c2w)
    c2w = c2w[start:end]
    intrins = intrins[start:end]
    
    # Convert camera-to-world to world-to-camera
    w2c = np.linalg.inv(c2w)
    
    # Add width and height to intrinsics
    N = len(w2c)
    intrins_full = np.zeros((N, 6), dtype=np.float32)
    intrins_full[:, :4] = intrins
    intrins_full[:, 4] = img_width
    intrins_full[:, 5] = img_height
    
    print(f"Loaded {N} VIPE camera poses")
    return w2c, intrins_full


def save_vipe_cameras_as_droid(output_dir, w2c, intrins_full):
    """
    Save VIPE cameras in DROID-SLAM format
    
    Args:
        output_dir: Output directory
        w2c: (N, 4, 4) world-to-camera matrices
        intrins_full: (N, 6) intrinsics [fx, fy, cx, cy, W, H]
    """
    os.makedirs(output_dir, exist_ok=True)
    
    # Extract parameters
    W, H = intrins_full[0, 4], intrins_full[0, 5]
    focal = intrins_full[:, :2].mean()
    
    print(f"Saving VIPE cameras to {output_dir}")
    print(f"Image size: {int(W)}x{int(H)}, focal: {focal:.2f}")
    
    # Save cameras.npz (main format used by Dyn-HaMR)
    np.savez(
        f"{output_dir}/cameras.npz",
        height=H,
        width=W,
        focal=focal,
        intrins=intrins_full[:, :4],
        w2c=w2c,
    )
    print(f"Saved cameras.npz with {len(w2c)} frames")


def run_vipe(video_path, vipe_dir, vipe_root=None, vipe_python=None, vipe_pipeline=None):
    """Run VIPE without shell activation, using the project venv by default.

    ``VIPE_PYTHON`` or ``vipe_python`` can select a separate interpreter.
    Input/output paths are resolved before changing the child working directory.
    """
    project_root = Path(__file__).resolve().parents[4]
    root = Path(vipe_root or project_root / "third_party" / "vipe").expanduser().resolve()
    video = Path(video_path).expanduser().absolute()
    output = Path(vipe_dir).expanduser().absolute()
    configured_python = vipe_python or os.environ.get("VIPE_PYTHON")
    project_python = project_root / ".venv" / "bin" / "python"
    python = Path(configured_python or (project_python if project_python.is_file() else sys.executable)).expanduser().absolute()
    if not video.is_file():
        raise FileNotFoundError(f"VIPE input video not found: {video}")
    if not (root / "vipe" / "cli" / "main.py").is_file():
        raise FileNotFoundError(f"VIPE source not found in {root}; set data.vipe_root")
    if not python.is_file():
        raise FileNotFoundError(f"VIPE interpreter not found: {python}; set VIPE_PYTHON")
    cmd = [str(python), "-m", "vipe.cli.main", "infer", str(video), "--output", str(output)]
    if vipe_pipeline:
        cmd.extend(["--pipeline", str(vipe_pipeline)])
    print(f"Running VIPE in {root}: {shlex.join(cmd)}", flush=True)
    out = subprocess.call(cmd, cwd=str(root))
    if out != 0:
        print(f"WARNING: VIPE failed with exit code {out}")
    else:
        print("VIPE completed successfully")
    return out


def preprocess_cameras(cfg, overwrite=False):
    if not overwrite and is_nonempty(cfg.sources.cameras):
        print(f"FOUND CAMERAS IN {cfg.sources.cameras}")
        return

    # Check if we should use VIPE instead of DROID-SLAM
    use_vipe = cfg.get("use_vipe", False)
    vipe_dir = cfg.get("vipe_dir", None)
    
    if use_vipe and vipe_dir is not None:
        # Check if VIPE results exist for this sequence
        vipe_pose_path = os.path.join(vipe_dir, "pose", f"{cfg.seq}.npz")
        vipe_intrins_path = os.path.join(vipe_dir, "intrinsics", f"{cfg.seq}.npz")
        
        # If VIPE results don't exist, run VIPE
        if not (os.path.exists(vipe_pose_path) and os.path.exists(vipe_intrins_path)):
            print(f"VIPE results not found for sequence '{cfg.seq}', running VIPE...")
            
            # The source location is independent of the output directory.
            vipe_root = cfg.get("vipe_root", None)
            video_path = cfg.get("src_path", None)
            
            if video_path is None or not os.path.exists(video_path):
                raise FileNotFoundError(
                    f"Video path not found: {video_path}\n"
                    f"Cannot run VIPE. Please provide a valid 'src_path' in your config."
                )
            
            # Run VIPE
            out = run_vipe(video_path, vipe_dir, vipe_root, cfg.get("vipe_python", None), cfg.get("vipe_pipeline", None))
            if out != 0:
                raise RuntimeError(
                    f"VIPE failed with exit code {out}\n"
                    f"Check the command and error above. Select an installed Python "
                    f"with VIPE_PYTHON or data.vipe_python; conda is not required."
                )
        
        # Load VIPE results (after potentially running VIPE)
        if not (os.path.exists(vipe_pose_path) and os.path.exists(vipe_intrins_path)):
            raise FileNotFoundError(
                f"VIPE results not found after execution:\n"
                f"  Pose: {vipe_pose_path}\n"
                f"  Intrinsics: {vipe_intrins_path}\n"
                f"VIPE may have failed silently. Please check VIPE logs."
            )
        
        print(f"USING VIPE CAMERAS FROM {vipe_dir}")
        img_dir = cfg.sources.images
        map_dir = cfg.sources.cameras
        
        # Get frame range
        subseqs, shot_idcs = split_frames_shots(cfg.sources.images, cfg.sources.shots)
        shot_idx = np.where(shot_idcs == cfg.shot_idx)[0][0]
        start, end = subseqs[shot_idx]
        
        if not cfg.split_cameras:
            # only run on specified segment within shot
            end = start + cfg.end_idx
            start = start + cfg.start_idx
        
        # Load and convert VIPE cameras
        w2c, intrins_full = load_vipe_cameras(vipe_dir, cfg.seq, img_dir, start, end)
        save_vipe_cameras_as_droid(map_dir, w2c, intrins_full)
        return

    # Default: use DROID-SLAM
    print(f"RUNNING SLAM ON {cfg.seq}")
    img_dir = cfg.sources.images
    map_dir = cfg.sources.cameras
    subseqs, shot_idcs = split_frames_shots(cfg.sources.images, cfg.sources.shots)
    print(shot_idcs, cfg.shot_idx, np.where(shot_idcs == cfg.shot_idx), cfg.sources.images, cfg.sources.shots)
    print(subseqs)
    shot_idx = np.where(shot_idcs == cfg.shot_idx)[0][0]
    # run on selected shot
    start, end = subseqs[shot_idx]
    if not cfg.split_cameras:
        # only run on specified segment within shot
        end = start + cfg.end_idx
        start = start + cfg.start_idx
    intrins_path = cfg.sources.get("intrins", None)
    if intrins_path is not None:
        intrins_path = check_intrins(cfg.type, cfg.root, intrins_path, cfg.seq, cfg.split)

    print('img_dir, map_dir, start, end, intrins_path', img_dir, map_dir, start, end, intrins_path)
    # raise ValueERRROR
    cmd = get_command(
        img_dir,
        map_dir,
        start=start,
        end=end,
        intrins_path=intrins_path,
        overwrite=overwrite,
    )
    print(cmd)
    gpu = os.environ.get("CUDA_VISIBLE_DEVICES", 0)
    out = subprocess.call(f"CUDA_VISIBLE_DEVICES={gpu} {cmd}", shell=True)
    assert out == 0, "SLAM FAILED"