# ==========================================
# 0. HUGGING FACE ZERO-GPU INITIALIZATION
# ==========================================
try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(fn):
            return fn


# ==========================================
# 1. IMPORTS
# ==========================================
import os
import re
import tempfile
import numpy as np
import whisper
import gradio as gr

from groq import Groq
from f5_tts.api import F5TTS


# ==========================================
# 2. CONFIGURATION
# ==========================================

# IMPORTANT:
# Set your Groq API key as an environment variable.
# Do NOT put your real API key directly in the code.
GROQ_API_KEY = os.environ.get("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")

# Your reference voice
REFERENCE_VOICE_FILE = "my_voice_fixed.wav"

# Text spoken in your reference voice file
REFERENCE_TEXT = (
    "Hi, my name is Rabbani. "
    "I am a software engineer at Gamana Solutions."
)

if not os.path.exists(REFERENCE_VOICE_FILE):
    raise FileNotFoundError(
        f"Reference voice file '{REFERENCE_VOICE_FILE}' not found."
    )


# ==========================================
# 3. INITIALIZE MODELS
# ==========================================

groq_client = None
whisper_model = None
f5tts = None


# ==========================================
# 4. GROQ CLIENT
# ==========================================

def get_groq_client():

    global groq_client

    if not GROQ_API_KEY:
        raise ValueError(
            "GROQ_API_KEY is not configured."
        )

    if groq_client is None:
        groq_client = Groq(
            api_key=GROQ_API_KEY
        )

    return groq_client


# ==========================================
# 5. REMOVE FILLER WORDS
# ==========================================

def remove_filler_words(text):

    if not text:
        return ""

    # Convert multiple spaces into one
    text = re.sub(r"\s+", " ", text).strip()

    filler_patterns = [

        r"\bumm+\b",
        r"\bum+\b",
        r"\buh+\b",
        r"\bah+\b",

        r"\byou know\b",
        r"\bi mean\b",

        r"\bactually\b",
        r"\bbasically\b",

        r"\blike\b",
        r"\bsort of\b",
        r"\bkind of\b",

        r"\bwell\b",
    ]

    for pattern in filler_patterns:
        text = re.sub(
            pattern,
            "",
            text,
            flags=re.IGNORECASE
        )

    # Remove repeated spaces
    text = re.sub(r"\s+", " ", text)

    # Fix spaces before punctuation
    text = re.sub(r"\s+([,.!?])", r"\1", text)

    return text.strip()


# ==========================================
# 6. GRAMMAR CORRECTION
# ==========================================

def correct_grammar(text):

    if not text:
        return ""

    try:

        client = get_groq_client()

        response = client.chat.completions.create(

            model="llama-3.3-70b-versatile",

            messages=[

                {
                    "role": "system",
                    "content": """
You are a speech-cleaning and grammar-correction system.

Your task is ONLY to clean the user's spoken sentence.

Rules:

1. Remove filler words.
2. Correct grammar.
3. Correct sentence structure.
4. Remove unnecessary repetitions.
5. Keep the original meaning.
6. Do not add new information.
7. Do not answer the question.
8. Do not explain the correction.
9. Return ONLY the corrected sentence.
10. Keep the sentence natural and conversational.

Example:

Input:
"Umm, I want to know like how Salesforce LWC works, you know?"

Output:
"I want to know how Salesforce LWC works."

Input:
"Actually umm I am working on Salesforce and I want to create ten records."

Output:
"I am working on Salesforce, and I want to create ten records."
"""
                },

                {
                    "role": "user",
                    "content": text
                }

            ],

            max_tokens=200,

            temperature=0.1,

            timeout=15
        )

        cleaned_text = (
            response
            .choices[0]
            .message
            .content
            .strip()
        )

        return cleaned_text

    except Exception as error:

        print(
            f"Grammar correction error: {error}"
        )

        # Fallback:
        # Return filler-cleaned text
        return remove_filler_words(text)


# ==========================================
# 7. WHISPER TRANSCRIPTION
# ==========================================

def transcribe_audio(audio_file):

    global whisper_model

    if whisper_model is None:

        print("Loading Whisper tiny model...")

        whisper_model = whisper.load_model(
            "tiny"
        )

    result = whisper_model.transcribe(
        audio_file,
        fp16=False
    )

    text = result["text"].strip()

    return text


# ==========================================
# 8. F5-TTS VOICE CLONING
# ==========================================

def clone_voice(text):

    global f5tts

    if not text:
        return None

    try:

        if f5tts is None:

            print("Loading F5-TTS...")

            f5tts = F5TTS()

        # Temporary output file
        with tempfile.TemporaryDirectory() as temp_dir:

            output_file = os.path.join(
                temp_dir,
                "generated.wav"
            )

            result = f5tts.infer(

                ref_file=REFERENCE_VOICE_FILE,

                ref_text=REFERENCE_TEXT,

                gen_text=text,

                file_wave=output_file
            )

            # ==================================
            # Get audio from F5-TTS
            # ==================================

            if isinstance(result, tuple):

                audio_data = result[0]
                sample_rate = result[1]

            else:

                audio_data = result
                sample_rate = 24000

            # ==================================
            # Convert Tensor → NumPy
            # ==================================

            if hasattr(
                audio_data,
                "cpu"
            ):

                audio_data = (
                    audio_data
                    .cpu()
                    .numpy()
                )

            elif hasattr(
                audio_data,
                "numpy"
            ):

                audio_data = (
                    audio_data
                    .numpy()
                )

            audio_data = np.asarray(
                audio_data
            )

            audio_data = np.squeeze(
                audio_data
            )

            # Make independent copy
            audio_data = audio_data.copy()

        # Temporary file is deleted here

        return (
            sample_rate,
            audio_data
        )

    except Exception as error:

        print(
            f"F5-TTS error: {error}"
        )

        return None


# ==========================================
# 9. MAIN VOICE PIPELINE
# ==========================================

@spaces.GPU
def voice_pipeline(audio_input):

    if audio_input is None:

        return (
            "Please record your voice.",
            "",
            None
        )

    try:

        # ==================================
        # STEP 1: Whisper
        # ==================================

        print(
            "\nSTEP 1: Transcribing audio..."
        )

        original_text = transcribe_audio(
            audio_input
        )

        print(
            "Original:",
            original_text
        )


        # ==================================
        # STEP 2: Remove filler words
        # ==================================

        print(
            "\nSTEP 2: Removing filler words..."
        )

        filler_removed_text = (
            remove_filler_words(
                original_text
            )
        )

        print(
            "Filler removed:",
            filler_removed_text
        )


        # ==================================
        # STEP 3: Grammar correction
        # ==================================

        print(
            "\nSTEP 3: Correcting grammar..."
        )

        corrected_text = (
            correct_grammar(
                filler_removed_text
            )
        )

        print(
            "Corrected:",
            corrected_text
        )


        # ==================================
        # STEP 4: Voice cloning
        # ==================================

        print(
            "\nSTEP 4: Cloning voice..."
        )

        generated_audio = clone_voice(
            corrected_text
        )


        if generated_audio is None:

            return (
                original_text,
                corrected_text,
                None
            )


        # ==================================
        # STEP 5: Return everything
        # ==================================

        return (
            original_text,
            corrected_text,
            generated_audio
        )

    except Exception as error:

        print(
            f"Pipeline error: {error}"
        )

        return (
            f"Error: {error}",
            "",
            None
        )


# ==========================================
# 10. GRADIO UI
# ==========================================

with gr.Blocks(
    title="Rabbani Voice Cleaner & Cloner"
) as interface:

    gr.Markdown(
        "# 🎙️ Voice Cleaner + Grammar Correction + Voice Cloning"
    )

    gr.Markdown(
        """
Speak naturally into the microphone.

The system will:

**Your Voice → Whisper → Remove Fillers → Correct Grammar → Clone Your Voice**
"""
    )


    # ======================================
    # INPUT
    # ======================================

    gr.Markdown(
        "## 🎤 Step 1 — Speak"
    )

    audio_input = gr.Audio(
        sources=["microphone"],
        type="filepath",
        label="Record Your Voice"
    )


    process_button = gr.Button(
        "Process Voice",
        variant="primary"
    )


    # ======================================
    # ORIGINAL TRANSCRIPTION
    # ======================================

    gr.Markdown(
        "## 📝 Step 2 — Whisper Transcription"
    )

    original_output = gr.Textbox(
        label="Original Speech",
        interactive=False
    )


    # ======================================
    # CORRECTED TEXT
    # ======================================

    gr.Markdown(
        "## ✨ Step 3 — Corrected Sentence"
    )

    corrected_output = gr.Textbox(
        label="Clean + Grammatically Correct Sentence",
        interactive=False
    )


    # ======================================
    # CLONED AUDIO
    # ======================================

    gr.Markdown(
        "## 🔊 Step 4 — Cloned Voice"
    )

    cloned_audio = gr.Audio(
        label="Your Cloned Voice",
        type="numpy",
        interactive=False
    )


    # ======================================
    # BUTTON EVENT
    # ======================================

    process_button.click(

        fn=voice_pipeline,

        inputs=[
            audio_input
        ],

        outputs=[
            original_output,
            corrected_output,
            cloned_audio
        ]
    )


# ==========================================
# 11. START APPLICATION
# ==========================================

interface.queue().launch()