import gradio as gr
from gradio import wasm_utils
from fastrtc import ReplyOnPause, AlgoOptions, SileroVadOptions, AdditionalOutputs, WebRTC, get_cloudflare_turn_credentials_async, get_cloudflare_turn_credentials #get_hf_turn_credentials, 
import os
from dotenv import load_dotenv
import time
import numpy as np
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from backend.tts import synthesize_text
from backend.asr import transcribe_audio, transcribe_typhoon
from backend.utils import preprocess_audio, is_valid_turn, preprocess_audio_simplified
from backend.main import stream_chat_response
import json
from pydub import AudioSegment
import ast
from backend.utils import get_device
if get_device() == "mps":
    load_dotenv(override=True)


phone_waiting_sound = AudioSegment.from_mp3("frontend/phone-ringing-24000.wav")
sound_samples = np.array(phone_waiting_sound.get_array_of_samples(), dtype=np.int16)
if phone_waiting_sound.channels > 1:
    sound_samples = sound_samples.reshape((-1, phone_waiting_sound.channels)).mean(axis=1)
sound_samples = sound_samples.astype(np.float32) / 32768.0  # Normalize to [-1, 
def startup(_):
    yield (phone_waiting_sound.frame_rate, sound_samples)
    STARTUP_MESSAGE = "สวัสดีค่ะ พลอย 1577Homeshopping ยินดีให้บริการค่ะ"
    yield from synthesize_text(STARTUP_MESSAGE)
    time.sleep(1)
    yield AdditionalOutputs([{"role": "assistant", "content": STARTUP_MESSAGE}])

custom_css = """
/* Overall Gradio page styling: hot pink background */
body {
/*    background-color: #ff69b4; /* Hot pink */
    margin: 0;
    padding: 0;
    font-family: sans-serif;}
/* Title styling */
h1 {
    color: #fff;
    text-shadow: 1px 1px 2px #ff85a2;
    font-size: 2.5em;
    margin-bottom: 20px;
    text-align: center;
}
/* Style the column holding the telephone interface */
.phone-column {
    max-width: 350px !important; /* Limit the width of the phone column */
    margin: 0 auto;              /* Center the column */
    border-radius: 20px;
    background-color: #f9cb9c;   /* Lighter pink for telephone interface */
    box-shadow: 0 0 15px rgba(0, 0, 0, 0.2);
    padding: 20px;
}
/* Conversation history box styling */
#conversation-history-chatbot {
    background-color: #f9cb9c;   /* Lighter pink for conversation history */
    border: 1px solid #ccc;
    border-radius: 10px;
    padding: 10px;
    box-shadow: 0 0 15px rgba(0, 0, 0, 0.2);
}
"""
def format_history_for_chatbot(history):
    """
    Formats the conversation history for display in the Gradio Chatbot.
    It creates readable strings for tool calls and tool results.
    """
    formatted_history = []
    if not history:
        return []
    for turn in history:
        role = turn.get("role")
        content = turn.get("content")
        tool_calls = turn.get("tool_calls")
        
        if role == "user":
            formatted_history.append({"role": "user", "content": content})
        elif role == "assistant":
            if tool_calls:
                # Display a user-friendly message for the tool call
                id = tool_calls[0]['id']
                func_name = tool_calls[0]['function']['name']
                func_args = tool_calls[0]['function']['arguments']
                display_content = f"<id>{id}</id><func_name>{func_name}</func_name><func_args>{func_args}</func_args>"
                # display_content = (
                #     f"**Calling Tool:**\n"
                #     f"```json\n"
                #     f"{{\n"
                #     f'  "name": "{func_name}",\n'
                #     f'  "arguments": {func_args}\n'
                #     f"}}\n"
                #     f"```"
                # )
                formatted_history.append({"role": "assistant", "content": display_content})
            else:
                # Regular assistant message
                formatted_history.append({"role": "assistant", "content": content})
        elif role == "tool":
            # Display a user-friendly message for the tool result
            id = turn.get("tool_call_id")
            result_content = json.dumps(json.loads(content), indent=2, ensure_ascii=False)
            display_content = f"<id>{id}</id><content>{content}</content>"
            # display_content = (
            #     f"**Tool Result:**\n"
            #     f"```json\n"
            #     f"{result_content}\n"
            #     f"```"
            # )
            # Represent tool results as if the "assistant" is providing them
            formatted_history.append({"role": "assistant", "content": display_content})
    
    return formatted_history
import re
def revert_to_openai_format(formatted_history):
    """
    Converts a history list formatted for the Gradio Chatbot UI back into
    the standard OpenAI API format. It parses custom string formats for
    tool calls and tool results.

    Args:
        formatted_history (list): A list of message dictionaries as they appear
                                  in the Gradio Chatbot component.

    Returns:
        list: A list of message dictionaries compliant with the OpenAI API format.
    """
    openai_history = []
    
    # Pre-compile regex patterns for efficiency
    # Pattern to find a tool call message
    tool_call_pattern = re.compile(
        r"<id>(.*?)</id><func_name>(.*?)</func_name><func_args>(.*?)</func_args>",
        re.DOTALL  # Use DOTALL in case arguments contain newlines
    )
    # Pattern to find a tool result message
    tool_result_pattern = re.compile(
        r"<id>(.*?)</id><content>(.*?)</content>",
        re.DOTALL
    )

    if not formatted_history:
        return []

    for turn in formatted_history:
        role = turn.get("role")
        content = turn.get("content") 
        # If content is None, treat it as an empty string for the regex search.
        if content is None:
            content = ""

        if role == "user":
            openai_history.append(turn)
            continue

        if role == "assistant":
            # Check if this is a formatted tool call
            tool_call_match = tool_call_pattern.search(content)
            if tool_call_match:
                call_id, func_name, func_args_str = tool_call_match.groups()
                
                # Reconstruct the original tool_calls structure
                reverted_turn = {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": call_id,
                            "type": "function",
                            "function": {
                                "name": func_name,
                                "arguments": func_args_str
                            },
                        }
                    ],
                }
                openai_history.append(reverted_turn)
                continue

            # Check if this is a formatted tool result (as per your formatter's logic)
            tool_result_match = tool_result_pattern.search(content)
            if tool_result_match:
                tool_call_id, tool_content = tool_result_match.groups()
                
                # Reconstruct the original tool message
                # NOTE: The role must be 'tool' for the API
                reverted_turn = {
                    "role": "tool",
                    "tool_call_id": tool_call_id,
                    "content": tool_content.strip() # Remove trailing space
                }
                openai_history.append(reverted_turn)
                continue
            
            # If no patterns match, it's a regular assistant message
            if turn.get("content") is not None:
                openai_history.append(turn)

    return openai_history


def response(audio: tuple[int, np.ndarray] | None, conversation_history):
    """
    Handles user audio input, transcribes it, streams LLM text via backend.main,
    and synthesizes chunks to audio while updating the conversation history.
    """
    print(f"WebRTC input SR: {audio[0]}")

    print(f"--- Latency Breakdown ---")
    # # stage_asr = "normal" #["normal,"gemini"]
    # print('-----------------------------')
    # print(f"Initial conver:{conversation_history}")
    # print('-----------------------------')

    conversation_history = revert_to_openai_format(conversation_history)
    # print(f"After convert:{conversation_history}")
    start_time = time.time()
    if conversation_history is None:
        conversation_history = []

    previous_history = list(conversation_history)

    if not audio or audio[1] is None or not np.any(audio[1]):
        print("No audio input detected; skipping response generation.")
        print(f"------------------------")
        return
    import soundfile as sf

    sample_rate, audio_array = audio

    try:
        sr , processed_audio = preprocess_audio_simplified((sample_rate, audio_array), target_sr=16000)
        
        print(sr, processed_audio.dtype, processed_audio.min(), processed_audio.max(), processed_audio.shape)

    except Exception as audio_err:
        print(f"Audio preprocessing failed: {audio_err}")
        print(f"------------------------")
        return
    
    silence_duration_s = 0.2

    # Calculate the number of samples corresponding to the silence duration
    silence_samples = int(16000 * silence_duration_s)
    
    # Create a silent audio segment (an array of zeros)
    # Ensure the dtype matches your processed audio for compatibility
    leading_silence = np.zeros(silence_samples, dtype=np.float32)

    # Prepend the silence to the beginning of your processed audio
    audio_with_padding = np.concatenate([leading_silence, processed_audio])
    
    print(f"Added {silence_duration_s}s of silence. New shape: {audio_with_padding.shape}")
    
    file_name = "temp.wav"
    sf.write(file_name, audio_with_padding, sr)
    t0 = time.time()
    transcription = transcribe_typhoon(file_name)
    # transcription = transcribe_audio( "debug_processed.wav")
    t_asr = time.time() - t0
    print(f"ASR:        {t_asr:.4f}s")

    if not transcription.strip():
        print("No valid transcription; skipping response generation.")
        print(f"------------------------")
        return
    user_turn = {"role": "user", "content": transcription}
    print(f"User: {transcription}")
    if is_valid_turn(user_turn):
        conversation_history.append(user_turn)
        yield AdditionalOutputs(format_history_for_chatbot(conversation_history))

    # print("Conversation history:", conversation_history)

    assistant_turn = {"role": "assistant", "content": ""}
    conversation_history.append(assistant_turn)
    
    # print(previous_history)
    history_for_stream = [dict(turn) for turn in previous_history if is_valid_turn(turn)]
    # print(f"history_for_stream{history_for_stream}")

    text_buffer = ""
    full_response = ""
    delimiter_count = 0
    n_threshold = 2
    max_n_threshold = 2
    lang = "th"
    chunk_count = 0
    first_chunk_sent = False
    start_llm_stream = time.time()

    try:
        for chunk in stream_chat_response(history_for_stream, transcription):
            # print(f"LLM chunk: {text_chunk}")
            if  isinstance(chunk, str):
                text_chunk = chunk

                i = 0
                while i < len(text_chunk):
                    char = text_chunk[i]
                    text_buffer += char
                    full_response += char

                    assistant_turn["content"] = full_response.strip()

                    is_delimiter = False
                    # if char in {' ', '\n'}:
                    if char == "|":
                        
                        #check the next character is number, not count as delimiter
                        # if i + 1 < len(text_chunk) and (text_chunk[i + 1].isdigit() or text_chunk[i - 1].isdigit()):
                        #     is_delimiter = False
                        # else:
                        # print(f"text_buffer before removing delimiter: '{text_buffer}'")
                        # text_buffer = text_buffer[:-1]  # Remove the delimiter from the buffer
                        # print(f"text_buffer after removing delimiter: '{text_buffer}'")
                        is_delimiter = True
                        
                        delimiter_count += 1
                        if i + 1 < len(text_chunk) and text_chunk[i + 1] == 'ๆ':
                            text_buffer += text_chunk[i + 1]
                            full_response += text_chunk[i + 1]
                            i += 1

                    send_now = False
                    if not first_chunk_sent:
                        if is_delimiter and text_buffer.strip():
                            send_now = True
                    else:
                        if delimiter_count >= n_threshold and text_buffer.strip():
                            send_now = True
                            if n_threshold < max_n_threshold:
                                n_threshold += 1

                    if send_now:
                        buffer_to_send = text_buffer.strip()
                        try:
                            if buffer_to_send and buffer_to_send.endswith('วันที่'):
                                buffer_to_send = buffer_to_send[:-len('วันที่')]
                            if buffer_to_send and first_chunk_sent and buffer_to_send.endswith('ค่ะ'):
                                buffer_to_send = buffer_to_send[:-len('ค่ะ')]
                        except Exception:
                            buffer_to_send = buffer_to_send.replace('ค่ะ', '')

                        if buffer_to_send:
                            chunk_count += 1
                            if chunk_count == 1:
                                first_llm_chunk_time = time.time()
                                t_llm_first_token = first_llm_chunk_time - start_llm_stream
                                print(f"LLM TTFC:   {t_llm_first_token:.4f}s (Time To First Chunk)")
                            yield from synthesize_text(buffer_to_send, lang=lang)
                            first_chunk_sent = True
                            text_buffer = ""
                            delimiter_count = 0
                            yield AdditionalOutputs(format_history_for_chatbot(conversation_history))

                    i += 1
            elif isinstance(chunk, dict) and "role" in chunk:
                # print(f"Received tool message for history: {chunk}")
                
                if chunk.get("content") is None:
                    chunk["content"] = ""
                conversation_history.insert(-1, chunk)

                # Update the chatbot UI to reflect the new history structure
                
                yield AdditionalOutputs(format_history_for_chatbot(conversation_history))
        
        if text_buffer.strip():
            buffer_to_send = text_buffer.strip()
            try:
                if buffer_to_send and buffer_to_send.endswith('วันที่'):
                    buffer_to_send = buffer_to_send[:-len('วันที่')]
                if buffer_to_send and first_chunk_sent and buffer_to_send.endswith('ค่ะ'):
                    buffer_to_send = buffer_to_send[:-len('ค่ะ')]
            except Exception:
                buffer_to_send = buffer_to_send.replace('ค่ะ', '')

            if buffer_to_send:
                chunk_count += 1
                if chunk_count == 1:
                    first_llm_chunk_time = time.time()
                    t_llm_first_token = first_llm_chunk_time - start_llm_stream
                    print(f"LLM TTFC:   {t_llm_first_token:.4f}s (Time To First Chunk)")
                yield from synthesize_text(buffer_to_send, lang=lang)
                first_chunk_sent = True
                text_buffer = ""
                delimiter_count = 0
                
                yield AdditionalOutputs(format_history_for_chatbot(conversation_history))

    except Exception as e:
        print(f"An error occurred during response generation or synthesis: {e}")
        error_message = "ขออภัยค่ะ เกิดข้อผิดพลาดบางอย่าง"
        try:
            yield from synthesize_text(error_message, lang=lang)
        except Exception as synth_error:
            print(f"Could not synthesize error message: {synth_error}")
        assistant_turn["content"] = (assistant_turn.get("content", "") + f" [Error: {e}]").strip()
        yield AdditionalOutputs(format_history_for_chatbot(conversation_history))

    total_latency = time.time() - start_time
    print(f"Total:      {total_latency:.4f}s")
    print(f"------------------------")



async def get_credentials():
    return await get_cloudflare_turn_credentials_async(hf_token=os.getenv('HF_TOKEN'))
   
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="orange", secondary_hue="orange")) as demo:
    gr.HTML("""<h1 style='text-align: center'>1577 Voicebot Demo</h1>""")
    with gr.Row():
        with gr.Column(scale=1, elem_classes=["phone-column"]):
            audio = WebRTC(
                mode="send-receive",
                modality="audio", 
                track_constraints={
                    "echoCancellation": True, 
                    "noiseSuppression": {"exact": True}, 
                    "autoGainControl": {"exact": True}
                },
                rtc_configuration=get_credentials,
                server_rtc_configuration=get_cloudflare_turn_credentials(ttl=360_000),
                icon="https://i.pinimg.com/originals/0c/67/5a/0c675a8e1061478d2b7b21b330093444.gif",
                icon_button_color="#17dbaa",
                pulse_color="#b0f83b",
                button_labels={"start": "Call", "stop": "Hang up", "waiting": "Connecting…"},
                icon_radius=45,
                height="650px",     
                width="100%",  
                container=False,    
                elem_id="phone-call-webrtc"
            )
        with gr.Column():
            conversation_history = gr.Chatbot(
                label="Conversation History", 
                type="messages", 
                value=[],
                height="675px",
                resizable=True,
                avatar_images=(None, "https://i.pinimg.com/originals/0c/67/5a/0c675a8e1061478d2b7b21b330093444.gif"),
            )
    gr.DeepLinkButton()

    audio.stream(
        fn=ReplyOnPause(
            response,
            algo_options=AlgoOptions(
                audio_chunk_duration=1.35,
                started_talking_threshold=0.35,
                speech_threshold=0.2
            ),
            model_options=SileroVadOptions(
                threshold=0.65,
                min_speech_duration_ms=200,
                max_speech_duration_s=float("inf"),
                min_silence_duration_ms=1200,
            ),
            can_interrupt=False,
            startup_fn=startup,
        ),
        inputs=[audio, conversation_history], 
        outputs=[audio],
        concurrency_limit=1000,
        time_limit=8192
    )

    audio.on_additional_outputs(
        lambda history: history,
        outputs=[conversation_history],
        queue=True,
        show_progress="hidden"
    )

demo.queue(default_concurrency_limit=10)
demo.launch(debug=True, show_error=True, share=True)