import gradio as gr
from transformers import pipeline
import re
import time

# Load both ASR models
whisper_asr = pipeline("automatic-speech-recognition", model="1morecupofhottea/Whisper-Code-Switching-Kh-En")
wav2vec_asr = pipeline("automatic-speech-recognition", model="1morecupofhottea/wav2vec2-Code-Switching-Kh_En")

def clean_transcript(text: str) -> str:
    # Remove <km>, </km>, <en>, </en>, etc.
    return re.sub(r"</?[^>]+>", "", text).strip()

def transcribe(audio, model_choice):
    if audio is None:
        return "Please upload or record an audio file first."
    
    try:
        if model_choice == "Whisper Model":
            result = whisper_asr(audio)
            return f"Transcription Complete:\n\n{result['text']}"
        else:
            result = wav2vec_asr(audio)
            cleaned_text = clean_transcript(result["text"])
            return f"Transcription Complete:\n\n{cleaned_text}"
    except Exception as e:
        return f"Error during transcription: {str(e)}"

def clear_all():
    return None, "Whisper Model", ""

# Custom CSS with dark theme
custom_css = """
/* 
DARK THEME COLOR PALETTE:
- Background: #000000 (Pure black - main background)
- Surface Dark: #1a1a1a (Dark gray - cards, containers)
- Surface Medium: #2d2d2d (Medium gray - inputs, components)
- Border: #404040 (Gray - borders, dividers)
- Text Primary: #ffffff (White - main text, headers)
- Text Secondary: #b0b0b0 (Light gray - secondary text)
- Accent: #4a9eff (Blue - buttons, highlights)
- Accent Hover: #357abd (Darker blue - hover states)
- Success: #28a745 (Green - success states)
- Warning: #ffc107 (Yellow - warnings)
*/

/* Global Styles */
.gradio-container {
    background: #000000 !important;
    font-family: 'Georgia', 'Times New Roman', serif !important;
    color: #ffffff !important;
    min-height: 100vh !important;
}

/* Header Styling */
.header-section {
    background: #1a1a1a;
    border: 2px solid #404040;
    border-radius: 12px;
    padding: 30px;
    margin-bottom: 25px;
    box-shadow: 0 4px 12px rgba(255, 255, 255, 0.1);
}

.header-section h1 {
    color: #ffffff !important;
    font-size: 2.4em !important;
    font-weight: 700 !important;
    text-align: center;
    margin-bottom: 15px;
    border-bottom: 3px solid #4a9eff;
    padding-bottom: 15px;
    text-shadow: 0 2px 4px rgba(0, 0, 0, 0.5);
}

.header-section p {
    color: #b0b0b0;
    font-size: 1.1em;
    text-align: center;
    margin: 8px 0;
    line-height: 1.6;
}

/* Main Content Cards */
.input-card, .output-card {
    background: #1a1a1a;
    border: 2px solid #404040;
    border-radius: 12px;
    padding: 25px;
    box-shadow: 0 4px 12px rgba(255, 255, 255, 0.05);
    margin-bottom: 20px;
}

/* Button Styling */
.primary-button {
    background: linear-gradient(135deg, #4a9eff, #357abd) !important;
    border: 2px solid #4a9eff !important;
    border-radius: 8px !important;
    padding: 14px 28px !important;
    font-size: 1.1em !important;
    font-weight: 600 !important;
    color: #ffffff !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 12px rgba(74, 158, 255, 0.3) !important;
}

.primary-button:hover {
    background: linear-gradient(135deg, #357abd, #2a5d8f) !important;
    border-color: #357abd !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 16px rgba(74, 158, 255, 0.4) !important;
}

.secondary-button {
    background: #2d2d2d !important;
    border: 2px solid #404040 !important;
    border-radius: 8px !important;
    padding: 12px 24px !important;
    font-size: 1em !important;
    font-weight: 500 !important;
    color: #ffffff !important;
    transition: all 0.3s ease !important;
}

.secondary-button:hover {
    background: #404040 !important;
    border-color: #555555 !important;
    transform: translateY(-1px) !important;
}

/* Component Styling */
.audio-component, .dropdown-component {
    border-radius: 8px !important;
    border: 2px solid #404040 !important;
    background: #2d2d2d !important;
    color: #ffffff !important;
}

/* Text Styling */
.output-text {
    background: #2d2d2d !important;
    border-radius: 8px !important;
    border: 2px solid #404040 !important;
    padding: 20px !important;
    font-size: 1.05em !important;
    line-height: 1.7 !important;
    font-family: 'Consolas', 'Courier New', monospace !important;
    color: #ffffff !important;
}

/* Features Section */
.features-section {
    background: #1a1a1a;
    border: 2px solid #404040;
    border-radius: 12px;
    padding: 25px;
    margin-top: 25px;
    box-shadow: 0 4px 12px rgba(255, 255, 255, 0.05);
}

.feature-item {
    margin-bottom: 15px;
    color: #b0b0b0;
    font-size: 1.05em;
    line-height: 1.6;
    padding: 10px;
    background: #2d2d2d;
    border-radius: 6px;
    border-left: 4px solid #4a9eff;
}

.feature-item strong {
    color: #ffffff;
}

/* Typography */
h3 {
    font-family: 'Georgia', 'Times New Roman', serif !important;
    color: #ffffff !important;
    font-weight: 600 !important;
    margin-bottom: 15px !important;
}

/* Component Overrides */
label {
    color: #ffffff !important;
    font-weight: 500 !important;
    font-family: 'Georgia', 'Times New Roman', serif !important;
}

::placeholder {
    color: #808080 !important;
    opacity: 0.8;
}

*:focus {
    outline: 2px solid #4a9eff !important;
    outline-offset: 2px !important;
}

/* Gradio Specific Overrides */
.gradio-textbox textarea,
.gradio-dropdown,
.gradio-audio {
    border: 2px solid #404040 !important;
    border-radius: 8px !important;
    background: #2d2d2d !important;
    color: #ffffff !important;
    font-family: inherit !important;
}

.gradio-textbox textarea:focus,
.gradio-dropdown:focus,
.gradio-audio:focus {
    border-color: #4a9eff !important;
    box-shadow: 0 0 0 3px rgba(74, 158, 255, 0.2) !important;
}

/* Scrollbar Styling */
::-webkit-scrollbar {
    width: 8px;
}

::-webkit-scrollbar-track {
    background: #1a1a1a;
}

::-webkit-scrollbar-thumb {
    background: #404040;
    border-radius: 4px;
}

::-webkit-scrollbar-thumb:hover {
    background: #4a9eff;
}

/* Selection Styling */
::selection {
    background: #4a9eff;
    color: #ffffff;
}

/* Responsive Design */
@media (max-width: 768px) {
    .header-section h1 {
        font-size: 2em !important;
    }
    
    .input-card, .output-card {
        padding: 20px;
        margin-bottom: 15px;
    }
    
    .primary-button, .secondary-button {
        padding: 12px 20px !important;
        font-size: 0.95em !important;
    }
}

/* Animation for better UX */
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(10px); }
    to { opacity: 1; transform: translateY(0); }
}

.input-card, .output-card, .features-section, .header-section {
    animation: fadeIn 0.5s ease-out;
}
"""

# Create the main interface
with gr.Blocks(css=custom_css, title="CS-ASR | Code-Switching Speech Recognition") as demo:
    # Header Section
    with gr.Column(elem_classes="header-section"):
        gr.HTML("""
            <h1>Code-Switching ASR Platform</h1>
            <p>Speech Recognition for Khmer-English Code-Switching</p>
        """)
    
    # Main Content
    with gr.Row():
        # Input Section
        with gr.Column(scale=1, elem_classes="input-card"):
            gr.HTML("<h3 style='color: #ffffff; margin-bottom: 20px; font-weight: 600;'>Audio Input</h3>")
            
            audio_input = gr.Audio(
                sources=["microphone", "upload"], 
                type="filepath", 
                label="Record or Upload Audio",
                elem_classes="audio-component"
            )
            
            model_selector = gr.Dropdown(
                choices=[
                    "Whisper Model", 
                    "Wav2Vec2 Model"
                ], 
                value="Whisper Model", 
                label="Select AI Model",
                elem_classes="dropdown-component",
                info="Choose the model that best fits your needs"
            )
            
            # Action Buttons
            with gr.Row():
                transcribe_button = gr.Button(
                    "Start Transcription", 
                    variant="primary",
                    elem_classes="primary-button",
                    scale=2
                )
                clear_button = gr.Button(
                    "Clear All", 
                    elem_classes="secondary-button",
                    scale=1
                )
        
        # Output Section  
        with gr.Column(scale=1, elem_classes="output-card"):
            gr.HTML("<h3 style='color: #ffffff; margin-bottom: 20px; font-weight: 600;'>Transcription Result</h3>")
            
            output_text = gr.Textbox(
                label="Your Transcription Will Appear Here",
                placeholder="Upload an audio file and click 'Start Transcription' to begin processing.",
                lines=12,
                elem_classes="output-text",
                interactive=False
            )
    
    # Event Handlers
    transcribe_button.click(
        fn=transcribe,
        inputs=[audio_input, model_selector],
        outputs=output_text,
        show_progress=True
    )
    
    clear_button.click(
        fn=clear_all,
        outputs=[audio_input, model_selector, output_text]
    )
    
    # Auto-transcribe when audio is uploaded (optional)
    audio_input.change(
        fn=lambda audio, model: transcribe(audio, model) if audio is not None else "",
        inputs=[audio_input, model_selector],
        outputs=output_text,
        show_progress=True
    )

# Launch with custom configuration
if __name__ == "__main__":
    demo.launch()