import sys
import types
_a = types.ModuleType('audioop')
_a.ratecv = lambda *a, **k: (b'', 0)
sys.modules['audioop'] = _a
sys.modules['pyaudioop'] = _a

import gradio as gr
import os
import time
import tempfile
import shutil
import subprocess
import numpy as np

print("ClearWave AI starting...")

# ── Services (flat structure - all files in root) ────────────────────
from denoiser    import Denoiser
from transcriber import Transcriber
from translator  import Translator

_denoiser    = Denoiser()
_transcriber = Transcriber()
_translator  = Translator()

# ── Config ───────────────────────────────────────────────────────────
INPUT_LANGS  = ["Auto Detect", "English", "Telugu", "Hindi", "Tamil", "Kannada"]
OUTPUT_LANGS = ["Telugu", "Hindi", "Tamil", "English", "Kannada"]
LANG_CODES   = {
    "Auto Detect": "auto",
    "English":     "en",
    "Telugu":      "te",
    "Hindi":       "hi",
    "Tamil":       "ta",
    "Kannada":     "kn",
}

# ── Pipeline ─────────────────────────────────────────────────────────
def process(audio_path, in_lang_label, out_lang_label, progress=gr.Progress()):
    if audio_path is None:
        return None, "Please upload audio.", "", "", "No audio"

    in_lang  = LANG_CODES.get(in_lang_label, "auto")
    out_lang = LANG_CODES.get(out_lang_label, "te")
    tmp      = tempfile.mkdtemp()
    t_total  = time.time()

    try:
        # Dept 1 — Denoise
        progress(0.1, desc="Dept 1: Denoising...")
        t0    = time.time()
        clean = _denoiser.process(audio_path, tmp)
        t1    = time.time() - t0

        # Dept 2 — Transcribe
        progress(0.4, desc="Dept 2: Transcribing...")
        t0                          = time.time()
        transcript, detected, tx_m  = _transcriber.transcribe(clean, in_lang)
        t2                          = time.time() - t0

        # Dept 3 — Translate
        progress(0.75, desc="Dept 3: Translating...")
        src          = detected if in_lang == "auto" else in_lang
        t0           = time.time()
        translated, tr_m = _translator.translate(transcript, src, out_lang)
        t3           = time.time() - t0

        total = time.time() - t_total
        progress(1.0, desc=f"Done in {total:.1f}s!")

        timing = (
            f"| Step | Time | Method |\n|---|---|---|\n"
            f"| Denoise    | {t1:.1f}s | noisereduce |\n"
            f"| Transcribe | {t2:.1f}s | {tx_m} |\n"
            f"| Translate  | {t3:.1f}s | {tr_m} |\n"
            f"| **Total**  | **{total:.1f}s** | |"
        )

        out_audio = os.path.join(tmp, "output.wav")
        shutil.copy(clean, out_audio)
        return out_audio, transcript, translated, timing, f"Done in {total:.1f}s"

    except Exception as e:
        import traceback
        return None, f"Error: {e}", "", traceback.format_exc(), "Failed"


# ── UI ───────────────────────────────────────────────────────────────
with gr.Blocks(title="ClearWave AI", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎵 ClearWave AI\n**Denoise · Transcribe · Translate**")

    with gr.Row():
        with gr.Column(scale=1):
            audio_in = gr.Audio(
                label="Upload Audio",
                type="filepath",
                sources=["upload", "microphone"],
            )
            in_lang  = gr.Dropdown(INPUT_LANGS,  value="Auto Detect", label="Input Language")
            out_lang = gr.Dropdown(OUTPUT_LANGS, value="Telugu",      label="Output Language")
            run_btn  = gr.Button("Process Audio", variant="primary", size="lg")
            status   = gr.Markdown("Upload audio and click Process.")

        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("Text"):
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("#### Transcript")
                            transcript_out = gr.Markdown("...")
                        with gr.Column():
                            gr.Markdown("#### Translation")
                            translation_out = gr.Markdown("...")
                with gr.Tab("Clean Audio"):
                    audio_out = gr.Audio(
                        label="Denoised",
                        type="filepath",
                        interactive=False,
                    )
                with gr.Tab("Timings"):
                    timing_out = gr.Markdown("...")

    run_btn.click(
        fn=process,
        inputs=[audio_in, in_lang, out_lang],
        outputs=[audio_out, transcript_out, translation_out, timing_out, status],
        show_progress=True,
        api_name=False,
    )

print("ClearWave AI ready!")
demo.launch()