import gradio as gr
import requests
import base64
import time
import json
import re

# ============================================================================
# MODAL BACKEND URLS - UPDATED
# ============================================================================
MODAL_TRANSCRIBE_URL = "https://vnmalleswar-payyavula--subtitle-generator-test-process-video.modal.run"
MODAL_BURN_URL = "https://vnmalleswar-payyavula--subtitle-generator-test-burn-edit-099d8b.modal.run"

# ============================================================================
# STEP 1: TRANSCRIBE VIDEO (Calls Modal Backend)
# ============================================================================

def transcribe_video(video_file, language="hindi", progress=gr.Progress()):
    """Call Modal backend to transcribe video - Returns video WITH subtitles burned"""
    
    if not video_file:
        return None, None, None, None, "❌ Please upload a video", ""
    
    try:
        progress(0.1, desc="Reading video file...")
        
        # Read and encode video
        with open(video_file, "rb") as f:
            video_bytes = f.read()
        
        video_base64 = base64.b64encode(video_bytes).decode('utf-8')
        
        progress(0.2, desc="Uploading to Modal server...")
        
        # Prepare request
        payload = {
            "video_base64": video_base64,
            "language": language,
            "position": "bottom-center"
        }
        
        progress(0.3, desc="Transcribing on GPU server (30-60 seconds)...")
        
        # Call Modal transcribe endpoint
        response = requests.post(
            MODAL_TRANSCRIBE_URL,
            json=payload,
            timeout=600
        )
        
        if response.status_code != 200:
            return None, None, None, None, f"❌ Server error: Status {response.status_code}\n{response.text}", ""
        
        result = response.json()
        
        # Check for errors
        if "error" in result:
            error_msg = result["error"]
            if "traceback" in result:
                error_msg += f"\n\n{result['traceback']}"
            return None, None, None, None, f"❌ Error: {error_msg}", ""
        
        progress(0.9, desc="Processing results...")
        
        # Get results
        srt_content = result.get("srt_content", "")
        video_base64_output = result.get("video_base64", "")  # Video WITH subtitles
        transcription = result.get("transcription", "")
        
        # DEBUG: Print first few lines of SRT
        print("=" * 60)
        print("SRT CONTENT RECEIVED:")
        print(srt_content[:500])
        print("=" * 60)
        
        # Save SRT file for download
        timestamp = int(time.time())
        srt_path = f"/tmp/subtitles_{timestamp}.srt"
        with open(srt_path, "w", encoding="utf-8") as f:
            f.write(srt_content)
        
        # Save original video (WITHOUT subtitles) for editing workflow
        original_video_path = f"/tmp/original_video_{timestamp}.mp4"
        with open(video_file, "rb") as f:
            original_bytes = f.read()
        with open(original_video_path, "wb") as f:
            f.write(original_bytes)
        
        # Save video WITH burned subtitles for preview
        preview_video_path = f"/tmp/preview_video_{timestamp}.mp4"
        if video_base64_output:
            preview_video_bytes = base64.b64decode(video_base64_output)
            with open(preview_video_path, "wb") as f:
                f.write(preview_video_bytes)
        else:
            preview_video_path = None
        
        progress(1.0, desc="Complete!")
        
        status = result.get("status", "Success")
        status += f"\n\n📊 **Statistics:**"
        status += f"\n- Original words: {result.get('original_word_count', 'N/A')}"
        status += f"\n- Subtitle words: {result.get('subtitle_word_count', 'N/A')}"
        status += f"\n- Words lost: {result.get('words_lost', 0)}"
        status += f"\n- Device: {result.get('device_used', 'N/A')}"
        status += f"\n- Model: {result.get('model_used', 'N/A')}"
        status += f"\n\n📝 **Transcription Preview:**\n{transcription[:300]}{'...' if len(transcription) > 300 else ''}"
        status += f"\n\n✅ **Video with subtitles is ready for preview!**"
        status += f"\n💡 You can edit subtitles below and re-burn if needed."
        
        return original_video_path, preview_video_path, srt_content, srt_path, status, ""
        
    except requests.Timeout:
        return None, None, None, None, "❌ Timeout: Processing took too long. Try a shorter video.", ""
    
    except Exception as e:
        import traceback
        error = f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
        return None, None, None, None, error, ""

# ============================================================================
# STEP 2: BURN EDITED SUBTITLES (Calls Modal Backend)
# ============================================================================

def burn_edited_subtitles(original_video_path, edited_srt_content, progress=gr.Progress()):
    """Call Modal backend to burn edited subtitles"""
    
    if not original_video_path:
        return None, "❌ No video found. Please transcribe a video first in Step 1."
    
    if not edited_srt_content or not edited_srt_content.strip():
        return None, "❌ No subtitles provided. Please edit subtitles first."
    
    try:
        progress(0.1, desc="Reading video file...")
        
        # Read original video
        with open(original_video_path, "rb") as f:
            video_bytes = f.read()
        
        video_base64 = base64.b64encode(video_bytes).decode('utf-8')
        
        progress(0.2, desc="Parsing edited subtitles...")
        
        # Parse edited SRT to subtitles array
        subtitles = parse_srt_string(edited_srt_content)
        
        print("=" * 60)
        print(f"PARSED {len(subtitles)} SUBTITLES:")
        for i, sub in enumerate(subtitles[:3]):
            print(f"Subtitle {i+1}: {sub}")
        print("=" * 60)
        
        if not subtitles:
            return None, "❌ Could not parse subtitles. Please check the format.\n\nMake sure each subtitle has:\n1. Index number\n2. Timestamp (HH:MM:SS,mmm --> HH:MM:SS,mmm)\n3. Text\n4. Blank line between subtitles"
        
        progress(0.3, desc="Uploading to Modal server...")
        
        # Prepare request
        payload = {
            "video_base64": video_base64,
            "subtitles": subtitles,
            "position": "bottom-center"
        }
        
        progress(0.4, desc="Burning subtitles on server (5-10 seconds)...")
        
        print(f"Sending {len(subtitles)} subtitles to Modal...")
        
        # Call Modal burn endpoint
        response = requests.post(
            MODAL_BURN_URL,
            json=payload,
            timeout=300
        )
        
        print(f"Response status: {response.status_code}")
        
        if response.status_code != 200:
            return None, f"❌ Server error: Status {response.status_code}\n{response.text}"
        
        result = response.json()
        
        # Check for errors
        if "error" in result:
            error_msg = result["error"]
            if "traceback" in result:
                error_msg += f"\n\n{result['traceback']}"
            return None, f"❌ Error: {error_msg}"
        
        progress(0.8, desc="Downloading processed video...")
        
        # Get output video
        video_base64_output = result.get("video_base64", "")
        
        if not video_base64_output:
            return None, "❌ No video returned from server"
        
        # Save final video
        timestamp = int(time.time())
        final_video_path = f"/tmp/final_video_{timestamp}.mp4"
        video_bytes = base64.b64decode(video_base64_output)
        with open(final_video_path, "wb") as f:
            f.write(video_bytes)
        
        print(f"Saved final video to {final_video_path}")
        
        progress(1.0, desc="Complete!")
        
        status = result.get("status", "Success")
        status += f"\n\n✅ Video ready! Burned {len(subtitles)} edited subtitles."
        status += f"\n\n🔍 **Debug Info:**"
        status += f"\n- Subtitles sent: {len(subtitles)}"
        status += f"\n- Server response: {result.get('subtitle_count', 'N/A')}"
        
        return final_video_path, status
        
    except requests.Timeout:
        return None, "❌ Timeout: Processing took too long."
    
    except Exception as e:
        import traceback
        error = f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
        return None, error

# ============================================================================
# HELPER FUNCTIONS - IMPROVED SRT PARSING
# ============================================================================

def parse_srt_string(srt_content):
    """Parse SRT string to subtitle array for Modal API - ROBUST VERSION"""
    subtitles = []
    
    # Split by double newline OR by subtitle index pattern
    blocks = re.split(r'\n\s*\n', srt_content.strip())
    
    for block in blocks:
        block = block.strip()
        if not block:
            continue
        
        lines = block.split('\n')
        
        # Need at least 3 lines: index, timestamp, text
        if len(lines) < 3:
            continue
        
        try:
            # Line 1: Index (must be a number)
            index_line = lines[0].strip()
            
            # Try to extract number from line
            index_match = re.search(r'(\d+)', index_line)
            if not index_match:
                continue
            
            index = int(index_match.group(1))
            
            # Line 2: Timestamp
            timestamp_line = lines[1].strip()
            
            # Parse timestamp: HH:MM:SS,mmm --> HH:MM:SS,mmm
            timestamp_match = re.search(
                r'(\d{2}:\d{2}:\d{2},\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2},\d{3})',
                timestamp_line
            )
            
            if not timestamp_match:
                print(f"⚠️ Could not parse timestamp: {timestamp_line}")
                continue
            
            start_time = timestamp_match.group(1)
            end_time = timestamp_match.group(2)
            
            # Lines 3+: Text (can be multiple lines)
            text = '\n'.join(lines[2:]).strip()
            
            if not text:
                print(f"⚠️ Empty text for subtitle {index}")
                continue
            
            subtitles.append({
                'index': index,
                'start': start_time,
                'end': end_time,
                'text': text
            })
            
        except Exception as e:
            print(f"⚠️ Error parsing block: {e}")
            print(f"Block content: {block[:100]}")
            continue
    
    return subtitles

# ============================================================================
# TEST PARSING FUNCTION
# ============================================================================

def test_srt_parsing(srt_content):
    """Test function to debug SRT parsing"""
    if not srt_content or not srt_content.strip():
        return "❌ No content to parse"
    
    subtitles = parse_srt_string(srt_content)
    
    result = f"✅ Parsed {len(subtitles)} subtitles\n\n"
    
    if len(subtitles) > 0:
        result += "📝 First 3 subtitles:\n\n"
        for i, sub in enumerate(subtitles[:3]):
            result += f"**Subtitle {sub['index']}:**\n"
            result += f"Time: {sub['start']} --> {sub['end']}\n"
            result += f"Text: {sub['text']}\n\n"
    else:
        result += "❌ No subtitles could be parsed!\n\n"
        result += "**Debug Info:**\n"
        result += f"Content length: {len(srt_content)} chars\n"
        result += f"First 200 chars:\n{srt_content[:200]}\n"
    
    return result

# ============================================================================
# GRADIO INTERFACE
# ============================================================================

css = """
#subtitle_editor {
    font-family: 'Courier New', monospace;
    font-size: 14px;
}
.modal-info {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 20px;
    border-radius: 10px;
    margin: 15px 0;
    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
}
.step-box {
    border: 2px solid #667eea;
    border-radius: 10px;
    padding: 20px;
    margin: 10px 0;
}
"""

with gr.Blocks(title="AI Subtitle Generator & Editor", css=css, theme=gr.themes.Soft()) as demo:
    
    # Header
    gr.Markdown(
        """
        # 🎬 AI Subtitle Generator & Editor
        
        ### Powered by Whisper Large V3 Turbo on Modal GPU
        
        **Two-Step Workflow:**
        1. 📝 **Transcribe** → Upload video, get auto-generated subtitles + preview video
        2. ✏️ **Edit & Re-burn** (Optional) → Modify subtitles and get updated video
        """
    )
    
    # Hidden states
    stored_video = gr.State(None)
    stored_subtitles_json = gr.State("")
    
    gr.Markdown("---")
    
    # ========================================================================
    # STEP 1: TRANSCRIBE
    # ========================================================================
    
    with gr.Group(elem_classes="step-box"):
        gr.Markdown("## 📝 Step 1: Upload & Transcribe Video")
        gr.Markdown("*This will automatically generate subtitles AND burn them into the video for preview*")
        
        with gr.Row():
            with gr.Column(scale=1):
                video_input = gr.Video(label="Upload Video File")
                language_dropdown = gr.Dropdown(
                    choices=["hindi", "english", "auto"],
                    value="hindi",
                    label="Audio Language",
                    info="Select the language spoken in the video"
                )
                transcribe_btn = gr.Button("🎙️ Transcribe & Generate Subtitles", variant="primary", size="lg")
            
            with gr.Column(scale=1):
                transcribe_status = gr.Textbox(
                    label="Transcription Status", 
                    lines=15,
                    interactive=False
                )
                srt_download = gr.File(
                    label="📥 Download SRT File",
                    interactive=False
                )
        
        # Preview video with subtitles
        with gr.Row():
            preview_video = gr.Video(
                label="📹 Preview: Video with Generated Subtitles",
                interactive=False
            )
        
        gr.Markdown(
            """
            <div class="modal-info">
            ⏱️ <b>Processing Time:</b> 30-60 seconds<br>
            🖥️ <b>Processing:</b> Modal GPU Server (T4)<br>
            ✨ <b>Output:</b> Video with burned subtitles + Editable SRT file
            </div>
            """
        )
    
    gr.Markdown("---")
    
    # ========================================================================
    # STEP 2: EDIT & RE-BURN (OPTIONAL)
    # ========================================================================
    
    with gr.Group(elem_classes="step-box"):
        gr.Markdown("## ✏️ Step 2: Edit Subtitles & Re-burn (Optional)")
        gr.Markdown("*If you're happy with the preview above, you're done! Otherwise, edit subtitles below and re-burn.*")
        
        with gr.Row():
            with gr.Column(scale=1):
                srt_editor = gr.Textbox(
                    label="✏️ Editable Subtitles (SRT Format)",
                    lines=20,
                    placeholder="Subtitles will appear here after transcription...\n\nYou can edit the text, timing, or delete entries.\n\nFormat:\n1\n00:00:00,000 --> 00:00:02,000\nSubtitle text\n\n2\n00:00:02,000 --> 00:00:04,000\nNext subtitle...",
                    elem_id="subtitle_editor"
                )
                
                with gr.Row():
                    test_parse_btn = gr.Button("🧪 Test Parsing", size="sm")
                    burn_btn = gr.Button("🔥 Re-burn with Edited Subtitles", variant="primary", size="lg", scale=2)
            
            with gr.Column(scale=1):
                parse_test_output = gr.Textbox(
                    label="🧪 Parse Test Results",
                    lines=10,
                    interactive=False,
                    visible=True
                )
                burn_status = gr.Textbox(
                    label="Re-burn Status",
                    lines=6,
                    interactive=False
                )
                final_video = gr.Video(
                    label="📹 Final Video with Edited Subtitles",
                    interactive=False
                )
        
        with gr.Accordion("💡 Editing Tips", open=False):
            gr.Markdown(
                """
                **How to edit subtitles:**
                - ✏️ **Change text:** Just modify the text lines
                - 🗑️ **Delete subtitle:** Remove entire block (index + timestamp + text)
                - ⏰ **Adjust timing:** Edit timestamp line (keep format: `HH:MM:SS,mmm --> HH:MM:SS,mmm`)
                - ➕ **Add subtitle:** Copy format and add new block
                - 🔢 **Re-index:** After deleting, you can renumber indices (1, 2, 3...)
                
                **Format rules:**
                ```
                1
                00:00:00,000 --> 00:00:02,000
                Your subtitle text here
                
                2
                00:00:02,000 --> 00:00:04,000
                Next subtitle text
                ```
                
                💡 **Use "Test Parsing" button to verify your edits before re-burning!**
                """
            )
        
        gr.Markdown(
            """
            <div class="modal-info">
            ⏱️ <b>Processing Time:</b> 5-10 seconds (no re-transcription!)<br>
            🖥️ <b>Processing:</b> Modal GPU Server<br>
            ⚡ <b>Speed:</b> Much faster than re-transcribing
            </div>
            """
        )
    
    gr.Markdown("---")
    
    # Footer
    with gr.Accordion("📚 More Information", open=False):
        gr.Markdown(
            f"""
            ### 🎯 Features:
            
            - ✅ Automatic subtitle generation using Whisper AI
            - ✅ Word-level timestamp accuracy
            - ✅ Instant video preview with subtitles
            - ✅ Edit subtitles without re-transcription
            - ✅ Support for Hindi, English, and auto-detection
            - ✅ Fast GPU processing on Modal
            
            ### 📊 Technical Details:
            
            | Component | Details |
            |-----------|---------|
            | Backend | Modal Serverless |
            | GPU | NVIDIA T4 |
            | Model | Whisper Large V3 Turbo |
            | Engine | faster-whisper |
            | Font | Arial, size 20 |
            | Position | Bottom center |
            | Max Words | 3 per subtitle |
            
            ### ⚡ Performance:
            
            | Operation | Time |
            |-----------|------|
            | Transcribe + Burn | 30-60 seconds |
            | Edit + Re-burn | 5-10 seconds |
            
            ### 🔧 Backend URLs:
            
            - Transcribe: `{MODAL_TRANSCRIBE_URL}`
            - Burn: `{MODAL_BURN_URL}`
            """
        )
    
    # ========================================================================
    # EVENT HANDLERS
    # ========================================================================
    
    # Step 1: Transcribe (returns video WITH subtitles)
    transcribe_btn.click(
        fn=transcribe_video,
        inputs=[video_input, language_dropdown],
        outputs=[stored_video, preview_video, srt_editor, srt_download, transcribe_status, stored_subtitles_json]
    )
    
    # Test SRT parsing
    test_parse_btn.click(
        fn=test_srt_parsing,
        inputs=[srt_editor],
        outputs=[parse_test_output]
    )
    
    # Step 2: Re-burn edited subtitles
    burn_btn.click(
        fn=burn_edited_subtitles,
        inputs=[stored_video, srt_editor],
        outputs=[final_video, burn_status]
    )

# ============================================================================
# LAUNCH
# ============================================================================

if __name__ == "__main__":
    demo.queue()
    demo.launch(
        share=False,
        show_error=True,
        server_name="0.0.0.0",
        server_port=7860
    )