import os
import torch
import gradio as gr
from transformers import pipeline
import transformers
import torch
import wave
import contextlib
import tempfile
from pydub import AudioSegment

MODEL_REPO_ID = os.environ["MODEL_REPO_ID"]
HF_TOKEN = os.environ["HF_TOKEN"]

device = 0 if torch.cuda.is_available() else -1
SEGMENT_LIMIT = 300  # 300s = 5 minutes
pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_REPO_ID,
    chunk_length_s=30,
    device=device,
    token=HF_TOKEN
)

def transcribe(file):
    if file is None:
        return "Please select an audio file."
    
    warning = ""
    
    try:
        all_results = []
        result = pipe(
            file, 
            batch_size=8, 
            generate_kwargs={
                "task": "transcribe", 
                "num_beams": 3
            },
        )
        
        if isinstance(result, dict) and 'chunks' in result:
            all_results.extend(result['chunks'])
        
        if all_results:
            timestamped = "\n".join([
                f"[{chunk['timestamp'][0]:.2f}:{chunk['timestamp'][1]:.2f}] {chunk['text'].strip()}" 
                for chunk in all_results if chunk['text'].strip()
            ])
            return timestamped + warning
        else:
            transcription = result.get('text', 'No transcription available') if isinstance(result, dict) else str(result)
            return transcription + warning
        
    except Exception as e:
        return f"Error during transcription: {str(e)}"

def transcribe_word_timestamps(file):
    if file is None:
        return "Please select an audio file.", "", ""
    
    if hasattr(file, 'name'):
        file_path = file.name
    else:
        file_path = file
    
    warning = ""

    try:
        all_chunks = []
        result = pipe(
            file_path, 
            generate_kwargs={
                "task": "transcribe", 
                "num_beams": 3,
                "condition_on_prev_tokens": False,
            },
            return_timestamps="word"
        )
        if isinstance(result, dict) and 'chunks' in result:
            all_chunks.extend(result['chunks'])
        
        if all_chunks:
            sample_text = ' '.join([chunk['text'].strip() for chunk in all_chunks[:3]])
            is_rtl = any('\u0600' <= char <= '\u06FF' for char in sample_text)
            
            import base64
            with open(file_path, 'rb') as audio_file:
                audio_data = audio_file.read()
                audio_base64 = base64.b64encode(audio_data).decode()
                
                # Determine MIME type based on file extension
                file_ext = file_path.lower().split('.')[-1]
                if file_ext in ['mp3', 'mpeg']:
                    audio_mime = "audio/mpeg"
                elif file_ext in ['wav']:
                    audio_mime = "audio/wav"
                elif file_ext in ['ogg']:
                    audio_mime = "audio/ogg"
                elif file_ext in ['m4a', 'aac']:
                    audio_mime = "audio/aac"
                else:
                    audio_mime = "audio/mpeg"  # Default fallback
                
                audio_url = f"data:{audio_mime};base64,{audio_base64}"
            
            
            html_content = f'''
                <div style="max-width: 800px; min-width: 800px; margin: 0 auto; padding: 20px; border: 1px solid #444; border-radius: 10px; background: #2d2d2d; color: #fff;">
                    <!-- Custom Audio Player -->
                    <div style="background: #1a1a1a; padding: 15px; border-radius: 8px; margin-bottom: 20px; box-shadow: 0 2px 8px rgba(0,0,0,0.3);">
                        <audio id="custom-audio-player" preload="metadata" style="width: 100%; margin-bottom: 10px;">
                            <source src="{audio_url}" type="{audio_mime}">
                            Your browser does not support the audio element.
                        </audio>
                        
                        <div style="display: flex; align-items: center; gap: 10px; margin-bottom: 10px;">
                            <button id="play-pause-btn" style="background: #0d6efd; color: white; border: none; padding: 10px 20px; border-radius: 5px; cursor: pointer; font-size: 16px; transition: all 0.2s; min-width: 55px;">
                                ▶️
                            </button>
                            <span id="time-display" style="font-family: monospace; font-size: 14px; color: #ccc;">
                                00:00 / 00:00
                            </span>
                        </div>
                    </div>
                    
                    <!-- Transcription with highlighting -->
                    <div style="background: #1a1a1a; padding: 15px; border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.3);">
                        <h3 style="margin-top: 0; color: #fff;">Transcription:</h3>
                        <div id="transcription-text" style="line-height: 1.8; padding: 15px; border: 1px solid #444; border-radius: 5px; direction: {"rtl" if is_rtl else "ltr"}; text-align: {"right" if is_rtl else "left"}; background: transparent; color: #fff;">
            '''

            for i, chunk in enumerate(all_chunks):
                word = chunk['text'].strip()
                start_time = chunk['timestamp'][0]
                end_time = chunk['timestamp'][1] if chunk['timestamp'][1] is not None else start_time + 0.5
                
                # Active highlight color fallback placeholder handled via plain inline styles safely
                html_content += f'<span class="word-span" data-word-index="{i}" data-start="{start_time}" data-end="{end_time}" style="cursor: pointer; padding: 4px 6px; margin: 2px; border-radius: 4px; display: inline-block; transition: all 0.2s; background: transparent; border: 1px solid transparent; color: #fff;">{word}</span> '
            
            html_content += '''
                    </div>
                    <div style="margin-top: 10px; font-size: 12px; color: #888;">
                        🎯 Click on any word to jump to that timestamp • Words will highlight as audio plays
                    </div>
                    <div style="margin-top: 5px; font-size: 12px; color: #888;">
                        ⏰ روی کلمات کلیک کن تا به زمان موردنظر بری • کلمات با صدا روشن می‌شوند
                    </div>
                </div>
            </div>
            '''
            
            import json
            words_json = json.dumps(all_chunks)
            
            return html_content, words_json, warning
        else:
            transcription = result.get('text', 'No transcription available') if isinstance(result, dict) else str(result)
            return f'<div style="min-height: 200px; color: #fff; background: #2d2d2d; padding: 20px; border-radius: 10px; border: 1px solid #444;">{transcription}</div>', "[]", warning
        
    except Exception as e:
        return f'<div style="min-height: 200px; color: #fff; background: #2d2d2d; padding: 20px; border-radius: 10px; border: 1px solid #444;">Error during word timestamp transcription: {str(e)}</div>', "[]", ""

basic_interface = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(sources=["upload", "microphone"], type="filepath", label="Upload or Record Audio")
    ],
    outputs=gr.Textbox(label="Transcription"),
    title="C1Tech Whisper Persian/فارسی",
    description="Upload an audio file or record directly. Outputs transcription."
)

with gr.Blocks(
    title="C1Tech Whisper Persian Transcription", 
    theme=gr.themes.Base(), 
    css="style.css"
) as demo:
    gr.Markdown("""
    # 🎙️ C1Tech Whisper Persian - Audio Transcription with Timestamps
    
    This application provides Persian speech-to-text transcription with precise timestamps.
    Choose the transcription type below. Both support file upload and microphone recording.
    
    </br>این برنامه گفتار فارسی را به متن تبدیل می‌کند و زمان دقیق هر بخش را هم مشخص می‌کند
    می‌توانید نوع تبدیل متن را از گزینه‌های زیر انتخاب کنید. هر دو گزینه از بارگذاری فایل و ضبط با میکروفون پشتیبانی می‌کنند
    """)
    
    with gr.Tab("🔤 Word-Level Timestamps"):
        with gr.Column():
            audio_input = gr.File(
                file_types=["audio"], 
                label="Upload Audio File",
                file_count="single"
            )
            
            transcribe_btn = gr.Button("Transcribe", variant="primary")
            
            word_level_output = gr.HTML(
                label="Audio Player with Word-Level Highlighting",
                elem_id="transcription_display",
                value='<div style="min-height: 40px; padding: 20px; text-align: center; color: #888; background: #2d2d2d; border: 1px solid #444; border-radius: 10px;">Upload an audio file and click "Transcribe" to see the interactive transcription with word-level timestamps.</div>'
            )
            words_data = gr.Textbox(visible=False, elem_id="words_data_hidden")
            warning_output = gr.Textbox(label="Warning", visible=False)

        def transcribe_and_setup(file):
            html, json_data, warning = transcribe_word_timestamps(file)
            return html, json_data, warning

        transcribe_btn.click(
            fn=transcribe_and_setup,
            inputs=[audio_input],
            outputs=[word_level_output, words_data, warning_output],
            js=f"""
            async function(file) {{
                return file;
            }}
            """
        ).then(
            fn=None,
            inputs=[word_level_output, words_data],
            outputs=None,
            js="""
            function(html_content, words_json) {
                if (!words_json || words_json === 'null' || words_json.trim() === '') {
                    return;
                }
                
                try {
                    var chunks = JSON.parse(words_json);
                    if (!chunks || !chunks.length) return;
                } catch(e) {
                    return;
                }
                
                setTimeout(function() {
                    var audio = document.getElementById('custom-audio-player');
                    var playPauseBtn = document.getElementById('play-pause-btn');
                    var timeDisplay = document.getElementById('time-display');
                    var wordSpans = document.querySelectorAll('.word-span');
                    
                    if (!audio || !playPauseBtn) return;
                    
                    var isPlaying = false;
                    var duration = 0;
                    
                    function formatTime(seconds) {
                        if (isNaN(seconds) || !isFinite(seconds)) return '00:00';
                        var mins = Math.floor(seconds / 60);
                        var secs = Math.floor(seconds % 60);
                        return mins.toString().padStart(2, '0') + ':' + secs.toString().padStart(2, '0');
                    }
                    
                    function updateDisplay() {
                        var currentTime = audio.currentTime || 0;
                        duration = audio.duration || 0;
                        if (timeDisplay) {
                            timeDisplay.textContent = formatTime(currentTime) + ' / ' + formatTime(duration);
                        }
                    }
                    
                    // Track actual play/pause states reliably directly from native events
                    audio.addEventListener('play', function() {
                        isPlaying = true;
                        playPauseBtn.textContent = '⏸️';
                    });
                    
                    audio.addEventListener('pause', function() {
                        isPlaying = false;
                        playPauseBtn.textContent = '▶️';
                    });
                    
                    audio.addEventListener('ended', function() {
                        isPlaying = false;
                        playPauseBtn.textContent = '▶️';
                    });
                    
                    playPauseBtn.addEventListener('click', function(e) {
                        e.preventDefault();
                        if (isPlaying) {
                            audio.pause();
                        } else {
                            audio.play().catch(function(err) { console.error(err); });
                        }
                    });
                    
                    audio.addEventListener('loadedmetadata', function() {
                        duration = audio.duration;
                        updateDisplay();
                    });
                    
                    audio.addEventListener('timeupdate', function() {
                        var currentTime = audio.currentTime;
                        updateDisplay();
                        
                        // Clear active word states back to baseline
                        for (var i = 0; i < wordSpans.length; i++) {
                            var span = wordSpans[i];
                            // Only touch non-hovered tokens to ensure style locks remain fluid
                            if (span.getAttribute('data-is-active') === 'true') {
                                span.style.backgroundColor = 'transparent';
                                span.style.color = '#fff';
                                span.style.fontWeight = 'normal';
                                span.style.transform = 'scale(1)';
                                span.style.borderColor = 'transparent';
                                span.removeAttribute('data-is-active');
                            }
                        }
                        
                        // Highlight current word match
                        for (var i = 0; i < wordSpans.length; i++) {
                            var span = wordSpans[i];
                            var start = parseFloat(span.getAttribute('data-start'));
                            var end = parseFloat(span.getAttribute('data-end'));
                            
                            if (currentTime >= start && currentTime <= end) {
                                span.style.backgroundColor = '#0d6efd';
                                span.style.color = '#fff';
                                span.style.fontWeight = 'bold';
                                span.style.transform = 'scale(1.05)';
                                span.style.borderColor = '#0d6efd';
                                span.setAttribute('data-is-active', 'true');
                                break;
                            }
                        }
                    });
                    
                    // Hook interaction handlers safely
                    for (var i = 0; i < wordSpans.length; i++) {
                        (function(index, span) {
                            span.addEventListener('click', function() {
                                var start = parseFloat(this.getAttribute('data-start'));
                                audio.currentTime = start;
                                audio.play().catch(function(err){});
                            });
                            
                            span.addEventListener('mouseenter', function() {
                                // Only apply hover treatment if it isn't the currently spoken active word
                                if (this.getAttribute('data-is-active') !== 'true') {
                                    this.style.backgroundColor = '#555';
                                    this.style.borderColor = '#666';
                                }
                            });
                            
                            span.addEventListener('mouseleave', function() {
                                // Gracefully fall back to seamless transparency if it's not active
                                if (this.getAttribute('data-is-active') !== 'true') {
                                    this.style.backgroundColor = 'transparent';
                                    this.style.borderColor = 'transparent';
                                }
                            });
                        })(i, wordSpans[i]);
                    }
                    
                    setTimeout(updateDisplay, 100);
                }, 300);
            }
            """
        )

    with gr.Tab("📝 Basic Transcription"):
        basic_interface.render()

if __name__ == "__main__":
    demo.queue().launch()