import os
import json
import subprocess
import uuid
import sys
import argparse
import shutil
import zipfile
from datetime import timedelta

def timestamp_to_srt(seconds):
    td = timedelta(seconds=seconds)
    total_seconds = int(td.total_seconds())
    micros = td.microseconds
    hours, remainder = divmod(total_seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    return f"{hours:02}:{minutes:02}:{seconds:02},{micros//1000:03}"

def json_to_srt(json_data):
    """
    Converts internal JSON subtitle format to SRT.
    If 'words' key is present, generates word-level timestamps (Karaoke/Editing style).
    Otherwise, uses segment-level timestamps.
    """
    srt_content = ""
    counter = 1
    
    for block in json_data:
        # Check if words detail is available for Word-Level SRT
        if isinstance(block, dict) and "words" in block and block["words"]:
            for word_obj in block["words"]:
                start = word_obj.get('start', 0)
                end = word_obj.get('end', 0)
                text = word_obj.get('word', "")
                
                srt_content += f"{counter}\n"
                srt_content += f"{timestamp_to_srt(start)} --> {timestamp_to_srt(end)}\n"
                srt_content += f"{text}\n\n"
                counter += 1
        else:
            # Fallback to segment level
            start = 0
            end = 0
            text = ""
            if isinstance(block, dict):
                start = block.get('start', 0)
                end = block.get('end', 0)
                text = block.get('text', "")
            elif isinstance(block, (list, tuple)) and len(block) >= 3:
                start, end, text = block[0], block[1], block[2]
                
            srt_content += f"{counter}\n"
            srt_content += f"{timestamp_to_srt(start)} --> {timestamp_to_srt(end)}\n"
            srt_content += f"{text}\n\n"
            counter += 1
            
    return srt_content



# ... (rest of the file until export_pack end)

    # 8. ZIP IT
    zip_filename = f"{export_name}.zip"
    zip_path = os.path.join(project_path, zip_filename)
    
    # Create zip from stage_dir (base_name is without extension)
    shutil.make_archive(os.path.join(project_path, export_name), 'zip', stage_dir)
    
    print(f"SUCCESS: Export Pack created at {zip_path}")
    
    # Cleanup
    try:
        shutil.rmtree(stage_dir)
    except: pass
    
    return zip_path




def get_video_dims(vid_path):
    """Returns (width, height, duration_frames)"""
    try:
        cmd_w = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width", "-of", "default=noprint_wrappers=1:nokey=1", vid_path]
        width = int(subprocess.check_output(cmd_w).decode().strip())
        
        cmd_h = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=height", "-of", "default=noprint_wrappers=1:nokey=1", vid_path]
        height = int(subprocess.check_output(cmd_h).decode().strip())
        
        cmd_dur = ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", vid_path]
        dur_sec = float(subprocess.check_output(cmd_dur).decode().strip())
        
        # Assume 30fps for calculation if not probed, but probing is better
        cmd_fps = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=r_frame_rate", "-of", "default=noprint_wrappers=1:nokey=1", vid_path]
        fps_str = subprocess.check_output(cmd_fps).decode().strip()
        num, den = map(int, fps_str.split('/'))
        fps = num / den if den > 0 else 30.0
        
        return width, height, int(dur_sec * fps), fps
    except Exception as e:
        print(f"Error probing video: {e}")
        return 1920, 1080, 300, 30.0

def render_segmented_overlays(ass_path, segments, video_path, output_dir):
    """
    Renders segments using a physical transparent PNG canvas to ensure alpha correctness.
    """
    width, height, _, fps = get_video_dims(video_path)
    ass_path_sanitized = ass_path.replace("\\", "/").replace(":", "\\:")
    
    # Generate Base Canvas (Robust Way)
    canvas_png = os.path.join(output_dir, "base_canvas.png")
    # FFmpeg create transparent png
    subprocess.run([
        "ffmpeg", "-y", "-f", "lavfi", "-i", f"color=c=black@0.0:s={width}x{height}", 
        "-frames:v", "1", "-c:v", "png", canvas_png
    ], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    
    overlay_data = []
    print(f"Rendering {len(segments)} subtitle segments (Mode: Canvas Overlay)...")
    
    for i, seg in enumerate(segments):
        start = seg.get('start', 0)
        end = seg.get('end', 0)
        duration = end - start
        if duration <= 0: continue
        
        filename = f"caption_{i}.mov"
        out_path = os.path.join(output_dir, filename)
        
        # Input is PNG LOOP (Infinite) -> Trim duration -> Apply ASS -> Encode PNG Codec
        cmd = [
            "ffmpeg", "-y",
            "-loop", "1", "-i", canvas_png,
            "-vf", f"format=rgba,setpts=PTS+{start}/TB,ass='{ass_path_sanitized}',setpts=PTS-{start}/TB,format=rgba",
            "-t", str(duration),
            "-c:v", "png",
            "-pix_fmt", "rgba",
            "-an",
            out_path
        ]
        
        try:
            subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
            rel_path = os.path.join("captions", filename).replace("\\", "/")
            overlay_data.append({ "path": rel_path, "start": start, "end": end, "index": i })
            print(f"  [Seg {i}] Rendered {duration:.2f}s")
        except subprocess.CalledProcessError as e:
            print(f"  [Seg {i}] Failed: {e}")
            
    # Cleanup canvas
    if os.path.exists(canvas_png): os.remove(canvas_png)

    return overlay_data

def create_premiere_xml(project_name, video_path, overlay_segments, duration_frames, width=1080, height=1920, timebase=30, video_file_id=None, audio_file_id=None, scale_value=100.0, face_data=None, source_width=1920, source_height=1080):
    """
    Generates a Premiere Pro XML (xmeml version 4) with segmented cuts and overlays.
    overlay_segments: List of dicts [{'path', 'start', 'end', 'index'}]
    """
    
    # Generate unique IDs
    def get_uid(): return str(uuid.uuid4())[:12]
    
    if not video_file_id: video_file_id = f"file-video-{get_uid()}"
    if not audio_file_id: audio_file_id = f"file-audio-{get_uid()}"
    
    sequence_uuid = str(uuid.uuid4())
    
    # helper for file blocks
    def get_file_block(fid, fpath, is_audio_only=False):
       audio_blk = "" if is_audio_only else "<audio><samplecharacteristics><depth>16</depth><samplerate>48000</samplerate></samplecharacteristics><channelcount>2</channelcount></audio>"
       width_f = source_width
       height_f = source_height
       return f"""<file id="{fid}"><name>{os.path.basename(fpath)}</name><pathurl>{fpath}</pathurl><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><duration>{duration_frames}</duration><media><video><samplecharacteristics><width>{width_f}</width><height>{height_f}</height><alpha>straight</alpha></samplecharacteristics></video>{audio_blk}</media></file>"""

    # --- VIDEO TRACK SEGMENTATION ---
    cuts = []
    
    frame_face_map = {}
    if face_data:
        s_factor = scale_value / 100.0
        src_cx = source_width / 2.0
        src_cy = source_height / 2.0
        
        for entry in face_data:
            frame_idx = entry.get('frame')
            faces = entry.get('faces', [])
            if not faces: continue
            
            best_face = max(faces, key=lambda f: (f[2]-f[0]) * (f[3]-f[1]))
            cx = (best_face[0] + best_face[2]) / 2.0
            cy = (best_face[1] + best_face[3]) / 2.0
            
            off_x = cx - src_cx
            off_y = cy - src_cy
            tgt_h = - (off_x * s_factor) / width
            tgt_v = - (off_y * s_factor) / height
            frame_face_map[frame_idx] = (tgt_h, tgt_v)

    fps_float = float(timebase)
    
    if overlay_segments:
        current_frame = 0
        last_cam_center = (0.0, 0.0)
        if 0 in frame_face_map: last_cam_center = frame_face_map[0]
        
        sorted_segs = sorted(overlay_segments, key=lambda x: x['start'])
        
        for seg in sorted_segs:
            start_f = int(seg['start'] * fps_float)
            end_f = int(seg['end'] * fps_float)
            
            if start_f > current_frame:
                cuts.append({
                    "start": current_frame, 
                    "end": start_f, 
                    "center": last_cam_center
                })
            
            # Determine Mode Center (Avoid Middle Split)
            candidates_h = []
            candidates_v = []
            for f_idx in range(start_f, end_f):
                if f_idx in frame_face_map:
                    pos = frame_face_map[f_idx]
                    candidates_h.append(round(pos[0], 2)) # Round to cluster
                    candidates_v.append(round(pos[1], 2))
            
            if candidates_h:
                import statistics
                try:
                    # MODE: Pick the most frequent position
                    # Multi-mode handling: min(multimode) ensures consistency
                    best_h = min(statistics.multimode(candidates_h))
                    best_v = min(statistics.multimode(candidates_v))
                    current_cam_center = (best_h, best_v)
                except:
                    current_cam_center = last_cam_center
            else:
                current_cam_center = last_cam_center

            cuts.append({"start": start_f, "end": end_f, "center": current_cam_center})
            last_cam_center = current_cam_center
            current_frame = end_f
            
        if current_frame < duration_frames:
            cuts.append({"start": current_frame, "end": duration_frames, "center": last_cam_center})
    else:
        cuts.append({"start": 0, "end": duration_frames, "center": (0.0, 0.0)})

    video_track_items = ""
    for cut in cuts:
        seg_start = cut['start']
        seg_end = cut['end']
        c_h, c_v = cut['center']
        if seg_end - seg_start <= 0: continue
        
        seg_id = f"clipitem-video-{get_uid()}"
        basic_motion = f"""<filter><effect><name>Basic Motion</name><effectid>basic</effectid><effectcategory>motion</effectcategory><effecttype>motion</effecttype><mediatype>video</mediatype><parameter authoringApp="PremierePro"><parameterid>scale</parameterid><name>Scale</name><value>{scale_value}</value></parameter><parameter authoringApp="PremierePro"><parameterid>center</parameterid><name>Center</name><value><horiz>{c_h:.5f}</horiz><vert>{c_v:.5f}</vert></value></parameter><parameter authoringApp="PremierePro"><parameterid>centerOffset</parameterid><name>Anchor Point</name><value><horiz>-0.5</horiz><vert>-0.5</vert></value></parameter></effect></filter>"""
        video_track_items += f"""<clipitem id="{seg_id}"><name>{os.path.basename(video_path)}</name><duration>{duration_frames}</duration><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><start>{seg_start}</start><end>{seg_end}</end><in>{seg_start}</in><out>{seg_end}</out>{get_file_block(video_file_id, video_path)}{basic_motion}</clipitem>"""

    # --- OVERLAY TRACK SEGMENTATION (Opus Style) ---
    track_overlay_block = ""
    if overlay_segments and len(overlay_segments) > 0:
        overlay_clips = ""
        fps_float = float(timebase) # Assuming roughly match
        
        for seg in overlay_segments:
            start_f = int(seg['start'] * fps_float)
            end_f = int(seg['end'] * fps_float)
            clip_dur = end_f - start_f
            if clip_dur <= 0: continue
            
            ov_name = seg['path']
            ov_fid = f"file-ov-{seg['index']}-{get_uid()}"
            ov_cid = f"clip-ov-{seg['index']}-{get_uid()}"
            
            # File block specifically for this WebP
            file_blk = f"""<file id="{ov_fid}"><name>{os.path.basename(seg['path'])}</name><pathurl>{seg['path']}</pathurl><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><duration>{clip_dur}</duration><media><video><samplecharacteristics><width>{width}</width><height>{height}</height><alpha>straight</alpha></samplecharacteristics></video></media></file>"""
            
            overlay_clips += f"""<clipitem id="{ov_cid}"><name>{os.path.basename(seg['path'])}</name><duration>{clip_dur}</duration><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><start>{start_f}</start><end>{end_f}</end><in>0</in><out>{clip_dur}</out>{file_blk}<compositemode>normal</compositemode></clipitem>"""
            
        track_overlay_block = f"<track>{overlay_clips}</track>"
    else:
        track_overlay_block = "<track></track>"

    # --- ASSEMBLE ---
    timecode_block = f"""<timecode><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><string>00:00:00:00</string><frame>0</frame><displayformat>NDF</displayformat></timecode>"""
    audio_blk = f"""<track><clipitem id="{audio_file_id}"><name>{os.path.basename(video_path)}</name><duration>{duration_frames}</duration><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><start>0</start><end>{duration_frames}</end>{get_file_block(video_file_id, video_path)}<sourcetrack><mediatype>audio</mediatype><trackindex>1</trackindex></sourcetrack></clipitem></track>"""

    return f"""<?xml version="1.0" encoding="UTF-8"?><xmeml version="4"><sequence id="{sequence_uuid}"><name>{project_name}_CutRef</name><duration>{duration_frames}</duration><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate>{timecode_block}<media><video><format><samplecharacteristics><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><width>{width}</width><height>{height}</height><pixelaspectratio>square</pixelaspectratio></samplecharacteristics></format><track>{video_track_items}</track>{track_overlay_block}</video><audio>{audio_blk}</audio></media></sequence></xmeml>"""

def export_pack(project_path, segment_index, output_format="premiere"):
    """
    Generates a ZIP Pack for the segment.
    """
    print(f"Starting Export Pack for Project: {os.path.basename(project_path)}, Segment: {segment_index}")
    
    # Paths
    proj_name = os.path.basename(project_path)
    cut_dir = os.path.join(project_path, "cuts")
    
    # 1. IDENTIFY VIDEO FILE
    video_file = None
    original_scale_file = None
    
    if os.path.exists(cut_dir):
        files = os.listdir(cut_dir)
        # Search for {index}_..._original_scale.mp4 or similar
        prefix_idx = f"{segment_index:03d}_"
        
        for f in files:
            if f.startswith(prefix_idx) and (f.endswith(".mp4") or f.endswith(".mov")):
                 video_file = os.path.join(cut_dir, f)
                 break
    
    if not video_file:
        print(f"Error: No video file found for segment {segment_index} in {cut_dir}")
        return
        
    print(f"Selected Video: {video_file}")

    # 2. IDENTIFY SUBTITLE FILES
    subs_dir = os.path.join(project_path, "subs_ass")
    ass_file = None
    
    if os.path.exists(subs_dir):
        sub_files = os.listdir(subs_dir)
        prefix_idx = f"{segment_index:03d}_"
        # Prioritize Clean Processed > Processed > Any
        patterns = [
            (lambda f: f.endswith(".ass") and f.startswith(prefix_idx) and "processed" in f and "original" not in f), 
            (lambda f: f.endswith(".ass") and f.startswith(prefix_idx) and "processed" in f), 
            (lambda f: f.endswith(".ass") and f.startswith(prefix_idx))
        ]
        for p in patterns:
            if ass_file: break
            for f in sub_files:
                if p(f):
                    ass_file = os.path.join(subs_dir, f)
                    break
    
    # JSON in 'subs' usually
    subs_json_dir = os.path.join(project_path, "subs")
    json_file = None
    if os.path.exists(subs_json_dir):
        sub_files = os.listdir(subs_json_dir)
        prefix_idx = f"{segment_index:03d}_"
        # Same pattern priority
        json_patterns = [
            (lambda f: f.endswith(".json") and f.startswith(prefix_idx) and "processed" in f),
            (lambda f: f.endswith(".json") and f.startswith(prefix_idx))
        ]
        for p in json_patterns:
            if json_file: break
            for f in sub_files:
                if p(f):
                    json_file = os.path.join(subs_json_dir, f)
                    break

    # 2.1 IDENTIFY FACE COORDS
    final_dir = os.path.join(project_path, "final")
    face_data = None
    if os.path.exists(final_dir):
        final_files = os.listdir(final_dir)
        prefix_idx = f"{segment_index:03d}_"
        for f in final_files:
            if f.startswith(prefix_idx) and f.endswith("_coords.json"):
                try:
                    with open(os.path.join(final_dir, f), 'r') as fd:
                        face_data = json.load(fd)
                        print(f"Found Face Coordinates: {f}")
                except Exception as e:
                    print(f"Face coords load error: {e}")
                break

    # 3. PREPARE STAGING
    export_name = f"export_{proj_name}_seg{segment_index}"
    stage_dir = os.path.join(project_path, export_name)
    
    if os.path.exists(stage_dir):
        try:
            shutil.rmtree(stage_dir)
        except Exception:
            import random
            stage_dir += f"_{random.randint(1000,9999)}"
            
    os.makedirs(stage_dir, exist_ok=True)
    
    # 4. COPY VIDEO
    dest_video = os.path.join(stage_dir, "video_cut.mp4")
    shutil.copy2(video_file, dest_video)
    
    # 5. RENDER OVERLAYS (SEGMENTED)
    overlay_segments = []
    if ass_file and json_file:
         try:
             with open(json_file, 'r', encoding='utf-8') as f:
                 jdata = json.load(f)
             
             # Extract segment list
             jdata_segs = []
             if isinstance(jdata, dict) and "segments" in jdata:
                 jdata_segs = jdata["segments"]
             elif isinstance(jdata, list):
                 jdata_segs = jdata
             
             if jdata_segs:
                 # Create 'captions' subfolder for organization
                 captions_dir = os.path.join(stage_dir, "captions")
                 os.makedirs(captions_dir, exist_ok=True)
                 
                 # Render into subfolder
                 overlay_segments = render_segmented_overlays(ass_file, jdata_segs, video_file, captions_dir)
             
         except Exception as e:
             print(f"Error preparing overlay segments: {e}")
    else:
        print("Missing ASS or JSON for subtitles. Skipping overlays.")

    # 6. GENERATE SRT (Standard)
    dest_srt = os.path.join(stage_dir, f"{proj_name}_Seg{segment_index}.srt")
    if json_file:
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                jdata_srt = json.load(f)
            if isinstance(jdata_srt, dict) and "segments" in jdata_srt:
                jdata_srt = jdata_srt["segments"]
            srt_content = json_to_srt(jdata_srt)
            with open(dest_srt, 'w', encoding='utf-8') as f:
                f.write(srt_content)
        except Exception: pass

    # 7. GENERATE XML
    width_src, height_src, frames, fps = get_video_dims(dest_video)
    
    # Validation for resolution mismatch (same as before)
    if face_data:
        max_x = 0
        for entry in face_data:
            for f in entry.get('faces', []):
                if len(f) >= 3 and f[2] > max_x: max_x = f[2]
        if max_x > width_src:
            print(f"Correction: Detecting 4K source based on face coords ({max_x} > {width_src})")
            width_src = 3840
            height_src = 2160
            
    print(f"Using Source Resolution: {width_src}x{height_src}")

    xml_content = create_premiere_xml(
        project_name=f"{proj_name}_Seg{segment_index}",
        video_path="video_cut.mp4", 
        overlay_segments=overlay_segments, # LIST PASSED HERE
        duration_frames=frames,
        width=1080, # Target Vertical Width
        height=1920,
        timebase=int(float(fps) + 0.5), 
        source_width=width_src,
        source_height=height_src
    )
    
    xml_output = os.path.join(stage_dir, "timeline.xml")
    with open(xml_output, "w", encoding="utf-8") as f:
        f.write(xml_content)
        
    print("Generated Custom Premiere XML (Opus-Style Segments).")

    # 8. ZIP IT
    zip_path = f"{stage_dir}.zip"
    shutil.make_archive(stage_dir, 'zip', stage_dir)
    
    print(f"SUCCESS: Export Pack created at {zip_path}")
    
    # Cleanup
    try:
        shutil.rmtree(stage_dir)
    except: pass
    
    return zip_path

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--project", required=True)
    parser.add_argument("--segment", type=int, required=True)
    parser.add_argument("--format", default="premiere")
    args = parser.parse_args()

    export_pack(args.project, args.segment, args.format)