backend/utils.py

import numpy as np
import librosa
import io
import os
import warnings
import tempfile
from pydub import AudioSegment
from dotenv import load_dotenv
from fastrtc import get_cloudflare_turn_credentials_async, get_cloudflare_turn_credentials

try:
    import torch
except ModuleNotFoundError:
    torch = None  # type: ignore

warnings.filterwarnings("ignore")



# --- Device Configuration ---
def get_device():
    """Gets the best available device for PyTorch."""
    if torch is None:
        return "cpu"
    if torch.cuda.is_available():
        return "cuda"
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "mps"
    else:
        return "cpu"
if get_device() == "cpu":
    load_dotenv(override=True)
if get_device() == "mps":
    load_dotenv(override=True)
device = get_device()
print(f"Using device: {device}")

# --- Cloud Credentials ---
async def get_async_credentials():
    """Asynchronously fetches Cloudflare TURN credentials."""
    return await get_cloudflare_turn_credentials_async(hf_token=os.getenv('HF_TOKEN'))

def get_sync_credentials(ttl=360_000):
    """Synchronously fetches Cloudflare TURN credentials."""
    return get_cloudflare_turn_credentials(ttl=ttl)

def setup_gcp_credentials():
    """Sets up Google Cloud credentials from an environment variable."""
    gcp_service_account_json_str = os.getenv("GCP_SERVICE_ACCOUNT_JSON")
    # print(gcp_service_account_json_str)
    if gcp_service_account_json_str:
        try:
            # Create a temporary file to store the credentials
            with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=".json") as temp_file:
                temp_file.write(gcp_service_account_json_str)
                gcp_credential_path = temp_file.name # Get the path to the temporary file
            # Set the environment variable that Google Cloud libraries expect
            os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gcp_credential_path
            print(f"Google Cloud credentials set from secret to: {gcp_credential_path}")
        except Exception as e:
            print(f"Error setting up Google Cloud credentials: {e}")
    else:
        print("Warning: GCP_SERVICE_ACCOUNT_JSON secret not found. Google Cloud services may fail.")
    # if gcp_service_account_json_str:
    #     print("GCP service account JSON loaded from environment variable.")
    # else:
    #     print("Warning: GCP_SERVICE_ACCOUNT_JSON is not set; Google Cloud clients may fail.")
    # return gcp_service_account_json_str
# --- Audio Processing ---
# def audiosegment_to_numpy(audio, target_sample_rate=16000):
#     samples = np.array(audio.get_array_of_samples(), dtype=np.float32)
#     if audio.channels > 1:
#         samples = samples.reshape((-1, audio.channels)).mean(axis=1)
#     if audio.frame_rate != target_sample_rate:
#         samples = librosa.resample(samples, orig_sr=audio.frame_rate, target_sr=target_sample_rate)
#     samples /= np.iinfo(audio.array_type).max
#     return samples

def audiosegment_to_numpy(audio, target_sample_rate=16000):
    """
    Convert pydub.AudioSegment to normalized numpy array in range [-1, 1].
    """
    samples = np.array(audio.get_array_of_samples(), dtype=np.float32)
    if audio.channels > 1:
        samples = samples.reshape((-1, audio.channels)).mean(axis=1)

    # Normalize to [-1, 1]
    samples /= np.iinfo(audio.array_type).max

    # Resample if needed
    if audio.frame_rate != target_sample_rate:
        samples = librosa.resample(samples, orig_sr=audio.frame_rate, target_sr=target_sample_rate)

    # Final safety normalization
    max_val = np.max(np.abs(samples))
    if max_val > 0:
        samples = samples / max_val

    return samples.astype(np.float32)


def preprocess_audio(audio, target_channels=1, target_sr=16000):
    """
    Ensures the audio is mono, target sample rate, and normalized to [-1, 1].
    
    Args:
        audio: tuple (sample_rate, audio_array)
    Returns:
        tuple: (target_frame_rate, normalized_audio)
    """
    target_frame_rate = target_sr
    
    sample_rate, audio_array = audio
    #save audio array for debug
    with open("debug_audio_array.npy", "wb") as f:
        np.save(f, audio_array)
    print(audio_array)
    print(audio_array[0])
    print(len(audio_array[0]))
    print(audio_array.dtype)

    # Convert to int16 PCM if needed
    # If input is already float, scale it correctly
    if audio_array.dtype != np.int16:
        audio_array = np.clip(audio_array, -1.0, 1.0)
        audio_array_int16 = (audio_array * 32767).astype(np.int16)
    else:
        audio_array_int16 = audio_array

    # Wrap as BytesIO for AudioSegment
    audio_bytes = audio_array_int16.tobytes()
    audio_io = io.BytesIO(audio_bytes)

    # Convert to AudioSegment
    segment = AudioSegment.from_raw(
        audio_io, 
        sample_width=2, 
        frame_rate=sample_rate, 
        channels=1
    )

    # Adjust channels & frame rate
    segment = segment.set_channels(target_channels)
    segment = segment.set_frame_rate(target_frame_rate)

    # Convert back to normalized numpy
    samples = audiosegment_to_numpy(segment, target_sample_rate=target_frame_rate)

    return (target_frame_rate, samples)

def preprocess_audio_simplified(audio, target_sr=16000):
    """
    Ensures the audio is mono, at the target sample rate, and normalized to [-1, 1].
    
    Args:
        audio: tuple (original_sr, audio_array)
               audio_array is a numpy array.
    Returns:
        tuple: (target_sr, normalized_audio)
    """
    original_sr, audio_array = audio

    # Ensure audio_array is float
    if audio_array.dtype not in [np.float32, np.float64]:
        # Normalize int16 or other int types to [-1, 1]
        audio_array = audio_array.astype(np.float32) / np.iinfo(audio_array.dtype).max
        
    # Ensure audio is mono
    # Assumes channels are in the first dimension if it's 2D
    if audio_array.ndim > 1 and audio_array.shape[0] > 1:
        audio_array = np.mean(audio_array, axis=0)
    # If shape is (1, N), flatten it to (N,)
    audio_array = audio_array.flatten()

    # Resample if needed
    if original_sr != target_sr:
        audio_array = librosa.resample(y=audio_array, orig_sr=original_sr, target_sr=target_sr)

    # Peak normalization
    max_val = np.max(np.abs(audio_array))
    if max_val > 0:
        audio_array = audio_array / max_val
        
    return (target_sr, audio_array.astype(np.float32))



def is_valid_turn(turn: dict) -> bool:
    """
    Checks if a conversation turn is valid for inclusion in the LLM history.
    A turn is valid if it has a role and meets role-specific criteria:
    - user: must have non-empty content.
    - assistant: must have EITHER non-empty content OR tool_calls.
    - tool: must have content and a tool_call_id.
    """
    if not isinstance(turn, dict) or "role" not in turn:
        return False

    role = turn.get("role")
    
    if role == "user":
        # User turn is valid only if it has non-empty text content.
        return bool(turn.get("content") and isinstance(turn.get("content"), str) and turn.get("content").strip())
        
    elif role == "assistant":
        # Assistant turn is valid if it has text content OR if it has tool_calls.
        has_content = bool(turn.get("content") and isinstance(turn.get("content"), str) and turn.get("content").strip())
        has_tool_calls = "tool_calls" in turn and turn["tool_calls"] is not None
        return has_content or has_tool_calls

    elif role == "tool":
        # Tool turn is valid if it has a tool_call_id and content.
        return "tool_call_id" in turn and "content" in turn

    # Reject any other roles or malformed turns.
    return False