import json
import random
import sys
import os # Ensure os is imported early
# Set LITELLM_MODE to PRODUCTION before litellm is imported
os.environ['LITELLM_MODE'] = 'PRODUCTION'
import threading
import time
import datetime
import uvicorn
import webbrowser
import tempfile
import asyncio
import re
import base64 # Ensure it's used if needed by OpenAIRealtimeHandler
from pathlib import Path
from typing import List, Dict, Any, AsyncGenerator, Optional

from loguru import logger
import litellm
import numpy as np
import webview
import platformdirs
import click # Added click
import openai # Ensure openai is imported for AsyncOpenAI

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import (
    HTMLResponse,
    StreamingResponse,
    JSONResponse,
    FileResponse,
)
from fastrtc import (
    AdditionalOutputs,
    ReplyOnPause,
    Stream,
    get_twilio_turn_credentials,
    AlgoOptions,
    SileroVadOptions,
    AsyncStreamHandler, # Added AsyncStreamHandler
    wait_for_item, # Added wait_for_item
)
from gradio.utils import get_space
from openai import OpenAI, AuthenticationError # OpenAI is already here, AuthenticationError too
from pydantic import BaseModel
from pydub import AudioSegment

# --- Import Configuration ---
# This 'settings' instance will be populated in main() and used throughout.
from .utils.config import (
    settings, 
    APP_VERSION, 
    OPENAI_TTS_PRICING, 
    OPENAI_TTS_VOICES, 
    AppSettings, # Added AppSettings for type hint
    OPENAI_REALTIME_MODELS, # Added for OpenAI backend model list
    OPENAI_REALTIME_VOICES, # Added for OpenAI backend
    OPENAI_REALTIME_PRICING, # Updated from OPENAI_REALTIME_PRICING_PER_MINUTE
)
# --- End Import Configuration ---


# Import raw environment variable accessors
# These will be used as defaults for argparse arguments
from .utils.env import (
    LLM_HOST_ENV,
    LLM_PORT_ENV,
    DEFAULT_LLM_MODEL_ENV,
    LLM_API_KEY_ENV,
    STT_HOST_ENV,
    STT_PORT_ENV,
    STT_MODEL_ENV,
    STT_LANGUAGE_ENV,
    STT_API_KEY_ENV,
    STT_NO_SPEECH_PROB_THRESHOLD_ENV,
    STT_AVG_LOGPROB_THRESHOLD_ENV,
    STT_MIN_WORDS_THRESHOLD_ENV,
    TTS_HOST_ENV,
    TTS_PORT_ENV,
    TTS_MODEL_ENV,
    DEFAULT_VOICE_TTS_ENV,
    TTS_API_KEY_ENV,
    DEFAULT_TTS_SPEED_ENV,
    TTS_ACRONYM_PRESERVE_LIST_ENV,
    APP_PORT_ENV,
    SYSTEM_MESSAGE_ENV,
    OPENAI_API_KEY_ENV, 
    OPENAI_REALTIME_MODEL_ENV, 
    OPENAI_REALTIME_VOICE_ENV,
)

# Import other utils functions
from .utils.tts import (
    get_voices,
    generate_tts_for_sentence,
    prepare_available_voices_data,
)
from .utils.llms import (
    get_models_and_costs_from_proxy,
    get_models_and_costs_from_litellm,
    calculate_llm_cost,
)
from .utils.misc import is_port_in_use
from .utils.stt import transcribe_audio, check_stt_confidence
from .utils.logging_config import setup_logging

# --- Global Variables (Runtime State - Not Configuration) ---
# These are primarily for managing the server and UI state, not app settings.
last_heartbeat_time: datetime.datetime | None = None
heartbeat_timeout: int = 15  # Seconds before assuming client disconnected
shutdown_event = threading.Event()  # Used to signal monitor thread to stop
pywebview_window = None  # To hold the pywebview window object if created
uvicorn_server = None # Global variable to hold the Uvicorn server instance
# --- End Global Configuration & State ---

# --- Constants ---
OPENAI_REALTIME_SAMPLE_RATE = 24000
# --- End Constants ---


# --- Chat History Saving Function ---
def save_chat_history(history: List[Dict[str, str]]):
    """Saves the current chat history to a JSON file named after the startup timestamp."""
    if not settings.chat_log_dir or not settings.startup_timestamp_str:
        logger.warning(
            "Chat log directory or startup timestamp not initialized. Cannot save history."
        )
        return

    log_file_path = settings.chat_log_dir / f"{settings.startup_timestamp_str}.json"
    logger.debug(f"Saving chat history to: {log_file_path}")
    try:
        with open(log_file_path, "w", encoding="utf-8") as f:
            json.dump(history, f, ensure_ascii=False, indent=2)
        logger.debug(f"Successfully saved chat history ({len(history)} messages).")
    except IOError as e:
        logger.error(f"Failed to save chat history to {log_file_path}: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred while saving chat history: {e}")


# --- Core Response Logic (Async Streaming with Background TTS - CLASSIC BACKEND) ---
async def response(
    audio: tuple[int, np.ndarray],
    chatbot: list[dict] | None = None,
) -> AsyncGenerator[Any, None]:
    """
    Handles audio input, performs STT, streams LLM response text chunks to UI,
    generates TTS concurrently, and yields final audio and updates.
    This function is used by the 'classic' backend.
    """
    # Access module-level variables set after arg parsing in main()
    # Ensure clients are initialized (should be, but good practice)
    if not settings.stt_client or not settings.tts_client:
        logger.error("STT or TTS client not initialized for classic backend. Cannot process request.")
        # Yield error state?
        return

    # Work with a copy to avoid modifying the input list directly and ensure clean state per call
    current_chatbot = (chatbot or []).copy()
    logger.info(
        f"--- Entering response function with history length: {len(current_chatbot)} ---"
    )
    # Extract only role and content for sending to the LLM API
    # Handle both old dict format and new ChatMessage model format during transition if needed,
    # but current_chatbot should ideally contain ChatMessage objects or dicts matching its structure.
    messages = []
    for item in current_chatbot:
        if isinstance(item, dict):
            messages.append({"role": item["role"], "content": item["content"]})
        elif hasattr(item, 'role') and hasattr(item, 'content'): # Check if it looks like ChatMessage
             messages.append({"role": item.role, "content": item.content})
        else:
            logger.warning(f"Skipping unexpected item in chatbot history: {item}")


    # Add system message if defined
    if settings.system_message:
        # Prepend system message if not already present (e.g., first turn)
        if not messages or messages[0].get("role") != "system":
            messages.insert(0, {"role": "system", "content": settings.system_message})
            logger.debug("Prepended system message to LLM input.")
        elif (
            messages[0].get("role") == "system"
            and messages[0].get("content") != settings.system_message
        ):
            # Update system message if it changed (though this shouldn't happen with current setup)
            messages[0]["content"] = settings.system_message
            logger.debug("Updated existing system message in LLM input.")

    # Signal STT processing start
    yield AdditionalOutputs(
        {
            "type": "status_update",
            "status": "stt_processing",
            "message": "Transcribing...",
        }
    )

    # --- Speech-to-Text using imported function ---
    stt_success, prompt, stt_response_obj, stt_error = await transcribe_audio(
        audio,
        settings.stt_client,
        settings.stt_model_arg, # Use the model name from initial args/env
        settings.current_stt_language,
        settings.stt_api_base,
    )

    if not stt_success:
        logger.error(f"STT failed: {stt_error}")
        stt_error_msg = {
            "role": "assistant",
            "content": f"[STT Error: {stt_error or 'Unknown STT failure'}]",
        }
        yield AdditionalOutputs({"type": "chatbot_update", "message": stt_error_msg})
        # Yield final state and status even on STT error to reset frontend
        logger.warning("Yielding final state after STT error...")
        yield AdditionalOutputs(
            {
                "type": "final_chatbot_state",
                "history": current_chatbot,
            }  # Yield original state
        )
        yield AdditionalOutputs(
            {
                "type": "status_update",
                "status": "idle",
                "message": "Ready (after STT error)",
            }
        )
        logger.info("--- Exiting response function after STT error ---")
        return

    # --- STT Confidence Check using imported function ---
    reject_transcription, rejection_reason = check_stt_confidence(
        stt_response_obj, # The full response object from STT
        prompt, # The transcribed text
        settings.stt_no_speech_prob_threshold,
        settings.stt_avg_logprob_threshold,
        settings.stt_min_words_threshold,
    )
    # Store relevant STT details for potential metadata logging
    stt_metadata_details = {}
    if hasattr(stt_response_obj, 'no_speech_prob'):
        stt_metadata_details['no_speech_prob'] = stt_response_obj.no_speech_prob
    if hasattr(stt_response_obj, 'avg_logprob'):
        stt_metadata_details['avg_logprob'] = stt_response_obj.avg_logprob
    # Add word count if needed: stt_metadata_details['word_count'] = len(prompt.split())

    if reject_transcription:
        logger.warning(f"STT confidence check failed: {rejection_reason}. Details: {stt_metadata_details}")
        # Yield status updates to go back to idle without processing this prompt
        yield AdditionalOutputs(
            {
                "type": "status_update",
                "status": "idle",
                "message": f"Listening (low confidence: {rejection_reason})",
            }
        )
        # Yield final state (unchanged history)
        yield AdditionalOutputs(
            {
                "type": "final_chatbot_state",
                "history": current_chatbot,  # Send back original history
            }
        )
        logger.info(
            "--- Exiting response function due to low STT confidence/word count ---"
        )
        return

    # --- Proceed if STT successful and confidence check passed ---
    # Create user message with metadata
    user_metadata = ChatMessageMetadata(
        timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(),
        stt_details=stt_metadata_details if stt_metadata_details else None # Add STT details if available
    )
    user_message = ChatMessage(role="user", content=prompt, metadata=user_metadata)

    # Add to the *copy* and yield update (convert to dict for frontend compatibility if needed)
    current_chatbot.append(user_message.model_dump()) # Store as dict in the list
    yield AdditionalOutputs({"type": "chatbot_update", "message": user_message.model_dump()})

    # Update messages list (for LLM) based on the modified copy
    messages.append({"role": user_message.role, "content": user_message.content})

    # Save history after adding user message (save_chat_history expects list of dicts)
    save_chat_history(current_chatbot)

    # --- Streaming Chat Completion & Concurrent TTS ---
    llm_response_stream = None
    full_response_text = ""
    sentence_buffer = "" # Buffer for accumulating raw text including tags
    last_tts_processed_pos = 0 # Tracks the index in the *cleaned* buffer processed for TTS
    final_usage_info = None
    llm_error_occurred = False
    first_chunk_yielded = False # Track if we yielded the first chunk for UI
    response_completed_normally = False  # Track normal completion
    total_tts_chars = 0  # Initialize TTS character counter
    tts_audio_file_paths: List[str] = [] # List to store FILENAMES of generated TTS audio files for this response

    try:
        # Signal waiting for LLM
        yield AdditionalOutputs(
            {
                "type": "status_update",
                "status": "llm_waiting",
                "message": "Waiting for AI...",
            }
        )
        llm_start_time = time.time()
        logger.info(
            f"Sending prompt to LLM ({settings.current_llm_model}) for streaming..."
        )
        llm_args_dict = {
            "model": settings.current_llm_model,
            "messages": messages, # Use the history including the user prompt
            "stream": True,
            "stream_options": {
                "include_usage": True
            },  # Request usage data in the stream
            **({"api_base": settings.llm_api_base} if settings.use_llm_proxy else {}),
        }
        if settings.llm_api_key: # This is for classic backend LLM
            llm_args_dict["api_key"] = settings.llm_api_key

        llm_response_stream = await litellm.acompletion(**llm_args_dict)


        async for chunk in llm_response_stream:
            logger.debug(f"CHUNK: {chunk}")
            # Check for content delta first
            delta_content = None
            if chunk.choices and chunk.choices[0].delta:
                delta_content = chunk.choices[0].delta.content

            if delta_content:
                # Yield text chunk immediately for UI update
                yield AdditionalOutputs(
                    {"type": "text_chunk_update", "content": delta_content}
                )
                first_chunk_yielded = True

                sentence_buffer += delta_content
                full_response_text += delta_content

            # Check for usage information in the chunk (often in the *last* chunk)
            # LiteLLM might attach it differently depending on the provider.
            # Let's check more robustly.
            chunk_usage = getattr(chunk, "usage", None) or getattr(
                chunk, "_usage", None
            )  # Check common attributes
            if chunk_usage:
                # If usage is already a dict, use it, otherwise try .dict()
                if isinstance(chunk_usage, dict):
                    final_usage_info = chunk_usage
                elif hasattr(chunk_usage, "dict"):
                    final_usage_info = chunk_usage.dict()
                else:
                    final_usage_info = vars(chunk_usage)  # Fallback to vars()

                logger.info(
                    f"Captured usage info from LLM chunk: {final_usage_info}"
                )

            # --- Process buffer for TTS ---
            # Clean the *entire* current buffer first, then process new parts ending in newline
            if delta_content: # Only process if there was new content
                buffer_cleaned_for_tts = re.sub(r"<think>.*?</think>", "", sentence_buffer, flags=re.DOTALL)
                new_cleaned_text = buffer_cleaned_for_tts[last_tts_processed_pos:]

                # Find the last newline in the newly added cleaned text
                split_pos = new_cleaned_text.rfind('\n')

                if split_pos != -1:
                    # Extract the chunk ready for TTS (up to and including the last newline)
                    tts_ready_chunk = new_cleaned_text[:split_pos + 1]

                    # Split this chunk into sentences
                    sentences_for_tts = tts_ready_chunk.split('\n')

                    for sentence_for_tts in sentences_for_tts:
                        sentence_for_tts = sentence_for_tts.strip()
                        if sentence_for_tts:
                            # Count characters, log, call generate_tts_for_sentence, handle audio path...
                            total_tts_chars += len(sentence_for_tts)
                            logger.debug(f"Generating TTS for sentence (cleaned): '{sentence_for_tts[:50]}...' ({len(sentence_for_tts)} chars)")
                            audio_file_path: Optional[str] = await generate_tts_for_sentence(
                                sentence_for_tts,
                                settings.tts_client, # Classic backend TTS client
                                settings.tts_model_arg, 
                                settings.current_tts_voice,
                                settings.current_tts_speed,
                                settings.tts_acronym_preserve_set,
                                settings.tts_audio_dir,
                            )
                            if audio_file_path:
                                tts_audio_file_paths.append(audio_file_path)
                                full_audio_file_path = settings.tts_audio_dir / audio_file_path
                                logger.debug(f"TTS audio saved to: {full_audio_file_path}")
                                try:
                                    audio_segment = await asyncio.to_thread(
                                        AudioSegment.from_file, full_audio_file_path, format="mp3"
                                    )
                                    sample_rate = audio_segment.frame_rate
                                    samples = np.array(audio_segment.get_array_of_samples()).astype(np.int16)
                                    logger.debug(f"Yielding audio chunk from file '{audio_file_path}' for sentence: '{sentence_for_tts[:50]}...'")
                                    yield (sample_rate, samples)
                                except Exception as read_e:
                                    logger.error(f"Failed to read/decode TTS audio file {full_audio_file_path}: {read_e}")
                            else:
                                logger.warning(f"TTS failed for sentence, skipping audio yield and file save: '{sentence_for_tts[:50]}...'")

                    # Update the position marker for the cleaned buffer
                    last_tts_processed_pos += len(tts_ready_chunk)


        # After the loop, process any remaining text in the cleaned buffer
        buffer_cleaned_for_tts = re.sub(r"<think>.*?</think>", "", sentence_buffer, flags=re.DOTALL)
        remaining_cleaned_text = buffer_cleaned_for_tts[last_tts_processed_pos:].strip()

        if remaining_cleaned_text:
            # Process the final remaining part for TTS
            total_tts_chars += len(remaining_cleaned_text)
            logger.debug(f"Generating TTS for remaining buffer (cleaned): '{remaining_cleaned_text[:50]}...' ({len(remaining_cleaned_text)} chars)")
            audio_file_path: Optional[str] = await generate_tts_for_sentence(
                remaining_cleaned_text,
                settings.tts_client, # Classic backend TTS client
                settings.tts_model_arg, 
                settings.current_tts_voice,
                settings.current_tts_speed,
                settings.tts_acronym_preserve_set,
                settings.tts_audio_dir,
            )
            if audio_file_path:
                tts_audio_file_paths.append(audio_file_path)
                full_audio_file_path = settings.tts_audio_dir / audio_file_path
                logger.debug(f"TTS audio saved to: {full_audio_file_path}")
                try:
                    audio_segment = await asyncio.to_thread(
                        AudioSegment.from_file, full_audio_file_path, format="mp3"
                    )
                    sample_rate = audio_segment.frame_rate
                    samples = np.array(audio_segment.get_array_of_samples()).astype(np.int16)
                    logger.debug(f"Yielding audio chunk from file '{audio_file_path}' for remaining buffer: '{remaining_cleaned_text[:50]}...'")
                    yield (sample_rate, samples)
                except Exception as read_e:
                    logger.error(f"Failed to read/decode TTS audio file {full_audio_file_path}: {read_e}")
            else:
                logger.warning(f"TTS failed for remaining buffer, skipping audio yield and file save: '{remaining_cleaned_text[:50]}...'")


        llm_end_time = time.time()
        logger.info(
            f"LLM streaming finished ({llm_end_time - llm_start_time:.2f}s). Full response length: {len(full_response_text)}"
        )
        logger.info(
            f"Total characters sent to TTS: {total_tts_chars}"
        )  # Log total TTS chars

        # --- Final Updates (After LLM stream and TTS generation/yielding) ---
        response_completed_normally = (
            not llm_error_occurred
        )  # Mark normal completion if no LLM error occurred

        # 1. Cost Calculation (LLM and TTS)
        cost_result = {}  # Initialize cost result dict
        tts_cost = 0.0  # Initialize TTS cost

        # Calculate TTS cost if applicable (Classic backend)
        if settings.is_openai_tts and total_tts_chars > 0:
            tts_model_used = settings.tts_model_arg 
            if tts_model_used in OPENAI_TTS_PRICING:
                price_per_million_chars = OPENAI_TTS_PRICING[tts_model_used]
                tts_cost = (total_tts_chars / 1_000_000) * price_per_million_chars
                logger.info(
                    f"Calculated OpenAI TTS cost for {total_tts_chars} chars ({tts_model_used}): ${tts_cost:.6f}"
                )
            else:
                logger.warning(
                    f"Cannot calculate TTS cost: Pricing unknown for model '{tts_model_used}'."
                )
        elif total_tts_chars > 0:
            logger.info(
                f"TTS cost calculation skipped (not using OpenAI TTS or 0 chars). Total chars: {total_tts_chars}"
            )

        cost_result["tts_cost"] = tts_cost  # Add TTS cost (even if 0) to the result

        # Calculate LLM cost (if usage info available - Classic backend)
        if final_usage_info:
            llm_cost_result = calculate_llm_cost(
                settings.current_llm_model, final_usage_info, settings.model_cost_data
            )
            # Merge LLM cost results into the main cost_result dict
            cost_result.update(llm_cost_result)
            logger.info("LLM cost calculation successful.")
        elif not llm_error_occurred:
            logger.warning(
                "No final usage information received from LLM stream, cannot calculate LLM cost accurately."
            )
            cost_result["error"] = "LLM usage info missing"
            cost_result["model"] = settings.current_llm_model
            # Ensure LLM cost fields are present but potentially zero or marked
            cost_result.setdefault("input_cost", 0.0)
            cost_result.setdefault("output_cost", 0.0)
            cost_result.setdefault(
                "total_cost", 0.0
            ) 

        # Yield combined cost update
        logger.info(f"Yielding combined cost update: {cost_result}")
        yield AdditionalOutputs({"type": "cost_update", "data": cost_result})
        logger.info("Cost update yielded.")

        # 2. Add Full Assistant Text Response to History (to the copy) with Metadata
        assistant_message_obj = None 
        if not llm_error_occurred and full_response_text:
            # Create assistant message metadata
            assistant_metadata = ChatMessageMetadata(
                timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(),
                llm_model=settings.current_llm_model,
                usage=final_usage_info, 
                cost=cost_result, 
                tts_audio_file_paths=tts_audio_file_paths if tts_audio_file_paths else None 
            )
            # Create the full message object
            assistant_message_obj = ChatMessage(
                role="assistant",
                content=full_response_text,
                metadata=assistant_metadata
            )
            assistant_message_dict = assistant_message_obj.model_dump() 

            if not current_chatbot or current_chatbot[-1] != assistant_message_dict:
                current_chatbot.append(assistant_message_dict)
                logger.info(
                    "Full assistant response (with metadata) added to chatbot history copy for next turn."
                )
                # Save history after adding assistant message
                save_chat_history(current_chatbot)
            else:
                logger.info(
                    "Full assistant response already present in history, skipping append."
                )
        elif not llm_error_occurred:
            logger.warning(
                "LLM stream completed but produced no text content. History not updated."
            )

        # 3. Yield Final Chatbot State Update (if response completed normally)
        if response_completed_normally:
            logger.info("Yielding final chatbot state update...")
            yield AdditionalOutputs(
                {
                    "type": "final_chatbot_state",
                    "history": current_chatbot,
                } 
            )
            logger.info("Final chatbot state update yielded.")

        # 4. Yield Final Status Update (always, should be the last yield in the success path)
        final_status_message = (
            "Ready" if response_completed_normally else "Ready (after error)"
        )
        logger.info(f"Yielding final status update ({final_status_message})...")
        yield AdditionalOutputs(
            {
                "type": "status_update",
                "status": "idle",
                "message": final_status_message,
            }
        )
        logger.info("Final status update yielded.")

    except AuthenticationError as e: 
        logger.error(f"OpenAI Authentication Error during classic backend processing: {e}")
        response_completed_normally = False
        llm_error_occurred = True 

        error_content = f"\n[Authentication Error: Check your API key ({e})]"
        if not first_chunk_yielded:
            llm_error_msg = {"role": "assistant", "content": error_content.strip()}
            yield AdditionalOutputs(
                {"type": "chatbot_update", "message": llm_error_msg}
            )
        else:
            yield AdditionalOutputs(
                {"type": "text_chunk_update", "content": error_content}
            )

        logger.warning("Yielding final chatbot state (after auth exception)...")
        yield AdditionalOutputs(
            {"type": "final_chatbot_state", "history": current_chatbot}
        )
        logger.warning("Final chatbot state (after auth exception) yielded.")

        logger.warning("Yielding final status update (idle, after auth exception)...")
        yield AdditionalOutputs(
            {
                "type": "status_update",
                "status": "idle",
                "message": "Ready (Auth Error)",
            }
        )
        logger.warning("Final status update (idle, after auth exception) yielded.")

    except Exception as e:
        logger.error(
            f"Error during LLM streaming or TTS processing (classic backend): {type(e).__name__} - {e}", exc_info=True
        )
        response_completed_normally = False 
        llm_error_occurred = True 

        error_content = f"\n[LLM/TTS Error: {type(e).__name__}]" 
        if not first_chunk_yielded:
            llm_error_msg = {"role": "assistant", "content": error_content.strip()}
            yield AdditionalOutputs(
                {"type": "chatbot_update", "message": llm_error_msg}
            )
        else:
            yield AdditionalOutputs(
                {"type": "text_chunk_update", "content": error_content}
            )

        logger.warning("Yielding final chatbot state (after exception)...")
        yield AdditionalOutputs(
            {
                "type": "final_chatbot_state",
                "history": current_chatbot,
            } 
        )
        logger.warning("Final chatbot state (after exception) yielded.")

        logger.warning("Yielding final status update (idle, after exception)...")
        yield AdditionalOutputs(
            {
                "type": "status_update",
                "status": "idle",
                "message": "Ready (after error)",
            }
        )
        logger.warning("Final status update (idle, after exception) yielded.")

    logger.info(
        f"--- Response function generator finished (Completed normally: {response_completed_normally}) ---"
    )


# --- Pydantic Models ---

class ChatMessageMetadata(BaseModel):
    """Optional metadata associated with a chat message."""
    timestamp: Optional[str] = None 
    llm_model: Optional[str] = None
    usage: Optional[Dict[str, Any]] = None # For classic: token counts; For OpenAI: token counts from events
    cost: Optional[Dict[str, Any]] = None # For classic: LLM/TTS cost; For OpenAI: output audio cost
    stt_details: Optional[Dict[str, Any]] = None 
    tts_audio_file_paths: Optional[List[str]] = None # Classic backend
    output_audio_duration_seconds: Optional[float] = None # OpenAI backend
    raw_openai_usage_events: Optional[List[Dict[str, Any]]] = None # For OpenAI backend to store usage events

class ChatMessage(BaseModel):
    """Represents a single message in the chat history."""
    role: str
    content: str
    metadata: Optional[ChatMessageMetadata] = None

# --- FastAPI Setup ---

class InputData(BaseModel):
    """Model for data received by the /input_hook endpoint."""
    webrtc_id: str
    chatbot: list[ChatMessage] 


# --- OpenAI Realtime Handler ---
class OpenAIRealtimeHandler(AsyncStreamHandler):
    def __init__(self, app_settings: "AppSettings") -> None:
        super().__init__(
            expected_layout="mono", 
            output_sample_rate=OPENAI_REALTIME_SAMPLE_RATE,
            input_sample_rate=OPENAI_REALTIME_SAMPLE_RATE, 
        )
        self.settings = app_settings
        self.connection = None
        self.output_queue = asyncio.Queue()
        self.client: Optional[openai.AsyncOpenAI] = None 
        self.current_stt_language = self.settings.current_stt_language
        self.current_openai_voice = self.settings.current_openai_voice

        # State for cost calculation (OpenAI Backend)
        self.current_output_audio_duration_seconds: float = 0.0
        self.current_input_tokens: int = 0
        self.current_output_tokens: int = 0
        self.raw_usage_events_for_turn: List[Dict[str, Any]] = []


    def copy(self):
        return OpenAIRealtimeHandler(self.settings)

    def _reset_turn_usage_state(self):
        self.current_output_audio_duration_seconds = 0.0
        self.current_input_tokens = 0
        self.current_output_tokens = 0
        self.raw_usage_events_for_turn = []
        logger.debug("OpenAIRealtimeHandler: Turn usage state reset.")

    async def start_up(self):
        logger.info("OpenAIRealtimeHandler: Starting up and connecting to OpenAI...")
        if not self.settings.openai_api_key:
            logger.error("OpenAIRealtimeHandler: OpenAI API Key not configured. Cannot connect.")
            await self.output_queue.put(AdditionalOutputs({"type": "status_update", "status": "error", "message": "OpenAI API Key missing."}))
            await self.output_queue.put(None) 
            return

        self.client = openai.AsyncOpenAI(api_key=self.settings.openai_api_key) 

        # Update language and voice if changed in settings
        if self.current_stt_language != self.settings.current_stt_language:
            self.current_stt_language = self.settings.current_stt_language
            logger.info(f"OpenAIRealtimeHandler: STT language updated to: {self.current_stt_language or 'auto-detect'}")
        
        if self.current_openai_voice != self.settings.current_openai_voice:
            self.current_openai_voice = self.settings.current_openai_voice
            logger.info(f"OpenAIRealtimeHandler: Output voice updated to: {self.current_openai_voice}")


        transcription_config: Dict[str, Any] = {"model": "whisper-1"}
        if self.current_stt_language: 
            transcription_config["language"] = self.current_stt_language
        
        # Parameters for session.update()
        session_params: Dict[str, Any] = {
            "turn_detection": {"type": "server_vad"},
            "input_audio_transcription": transcription_config,
            # "output_audio_generation" and voice selection removed as it causes an error
            # with the default OpenAI Realtime model (e.g., gpt-4o-mini-realtime-preview-2024-12-17).
            # The API will likely use a default voice or one associated with the account/model.
        }
        
        if self.current_openai_voice:
            logger.warning(
                f"OpenAIRealtimeHandler: A voice ('{self.current_openai_voice}') is configured for the OpenAI backend, "
                "but the current API model/version might not support explicit voice selection via client parameters. "
                "The default voice for the model will likely be used."
            )
        
        # Parameters for client.beta.realtime.connect()
        connect_kwargs: Dict[str, Any] = {
            "model": self.settings.openai_realtime_model_arg
        }
        # Voice parameter moved to session_params
        
        try:
            self._reset_turn_usage_state() # Reset usage for new connection/session
            async with self.client.beta.realtime.connect(**connect_kwargs) as conn:
                await conn.session.update(session=session_params) 
                self.connection = conn
                logger.info(f"OpenAIRealtimeHandler: Connection established with model {self.settings.openai_realtime_model_arg}, voice {self.current_openai_voice or 'default via API'}.")
                
                async for event in self.connection:
                    logger.debug(f"OpenAIRealtime Event: {event.type}")
                    if event.type == "input_audio_buffer.speech_started":
                        self.clear_queue() 
                        self._reset_turn_usage_state() # Reset for a new turn of speech
                        await self.output_queue.put(
                            AdditionalOutputs(
                                {
                                    "type": "status_update",
                                    "status": "stt_processing", 
                                    "message": "Listening...",
                                }
                            )
                        )
                    elif event.type == "conversation.item.input_audio_transcription.completed":
                        user_message = ChatMessage(
                            role="user", 
                            content=event.transcript, 
                            metadata=ChatMessageMetadata(timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat())
                        )
                        await self.output_queue.put(
                            AdditionalOutputs({"type": "chatbot_update", "message": user_message.model_dump()})
                        )
                        await self.output_queue.put(
                            AdditionalOutputs(
                                {
                                    "type": "status_update",
                                    "status": "llm_waiting", 
                                    "message": "AI Responding...",
                                }
                            )
                        )
                    elif event.type == "response.audio_transcript.done":
                        # --- Cost Calculation & Metadata (OpenAI Backend - Token Based) ---
                        input_cost = 0.0
                        output_cost = 0.0
                        total_cost = 0.0
                        model_name_for_pricing = self.settings.openai_realtime_model_arg
                        model_prices = OPENAI_REALTIME_PRICING.get(model_name_for_pricing)

                        if not model_prices:
                            # Try to find a base model name match by checking if the model_arg starts with a key
                            for base_model_key in OPENAI_REALTIME_PRICING.keys():
                                if model_name_for_pricing.startswith(base_model_key):
                                    model_prices = OPENAI_REALTIME_PRICING[base_model_key]
                                    logger.info(f"OpenAIRealtimeHandler: Found pricing for '{model_name_for_pricing}' using base model key '{base_model_key}'.")
                                    break
                        
                        if model_prices:
                            price_input_per_mil = model_prices.get("input", 0.0)
                            price_output_per_mil = model_prices.get("output", 0.0)
                            # price_cached_input_per_mil = model_prices.get("cached_input", 0.0) # Not currently used

                            input_cost = (self.current_input_tokens / 1_000_000) * price_input_per_mil
                            output_cost = (self.current_output_tokens / 1_000_000) * price_output_per_mil
                            total_cost = input_cost + output_cost
                            logger.info(
                                f"OpenAIRealtime Token Costs: Input Tokens: {self.current_input_tokens}, Output Tokens: {self.current_output_tokens}. "
                                f"Input Cost: ${input_cost:.6f}, Output Cost: ${output_cost:.6f}, Total: ${total_cost:.6f}"
                            )
                        else:
                            logger.warning(
                                f"OpenAIRealtimeHandler: Could not find pricing information for model '{model_name_for_pricing}' in OPENAI_REALTIME_PRICING. Costs will be reported as $0."
                            )

                        cost_data = {
                            "input_cost": input_cost,
                            "output_cost": output_cost,
                            "total_cost": total_cost,
                            "input_tokens": self.current_input_tokens,
                            "output_tokens": self.current_output_tokens,
                            "model": self.settings.openai_realtime_model_arg,
                            "output_audio_duration_seconds": round(self.current_output_audio_duration_seconds, 2), # Informational
                            "note": "Costs are token-based. Audio duration is informational."
                        }
                        await self.output_queue.put(AdditionalOutputs({"type": "cost_update", "data": cost_data}))
                        # --- End Cost Calculation ---

                        assistant_metadata = ChatMessageMetadata(
                            timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), 
                            llm_model=self.settings.openai_realtime_model_arg,
                            cost=cost_data,
                            output_audio_duration_seconds=round(self.current_output_audio_duration_seconds, 2),
                            usage={"input_tokens": self.current_input_tokens, "output_tokens": self.current_output_tokens},
                            raw_openai_usage_events=self.raw_usage_events_for_turn.copy()
                        )
                        assistant_message = ChatMessage(
                            role="assistant", 
                            content=event.transcript, 
                            metadata=assistant_metadata
                        )
                        await self.output_queue.put(
                            AdditionalOutputs({"type": "chatbot_update", "message": assistant_message.model_dump()})
                        )
                        await self.output_queue.put(
                            AdditionalOutputs(
                                {
                                    "type": "status_update",
                                    "status": "idle",
                                    "message": "Ready",
                                }
                            )
                        )
                        # Removed problematic final_chatbot_state with empty history
                        # The UI should append messages as they come. 'idle' signals end of turn.

                    elif event.type == "response.audio.delta":
                        audio_data_bytes = base64.b64decode(event.delta)
                        audio_data_np = np.frombuffer(audio_data_bytes, dtype=np.int16)
                        
                        num_samples = len(audio_data_np)
                        duration_seconds_chunk = num_samples / OPENAI_REALTIME_SAMPLE_RATE
                        self.current_output_audio_duration_seconds += duration_seconds_chunk

                        if audio_data_np.ndim == 1: 
                            audio_data_np = audio_data_np.reshape(1, -1)
                        await self.output_queue.put(
                            (
                                OPENAI_REALTIME_SAMPLE_RATE,
                                audio_data_np,
                            ),
                        )
                    elif event.type == "conversation.item.usage.completed" or event.type == "response.usage.completed":
                        logger.info(f"OpenAIRealtime Usage Event: {event.type} - Raw Usage Data: {event.usage}") # More detailed log
                        current_event_input_tokens = 0
                        current_event_output_tokens = 0

                        if hasattr(event.usage, 'input_tokens') and event.usage.input_tokens is not None:
                            current_event_input_tokens = event.usage.input_tokens
                            self.current_input_tokens += current_event_input_tokens
                        if hasattr(event.usage, 'output_tokens') and event.usage.output_tokens is not None:
                            current_event_output_tokens = event.usage.output_tokens
                            self.current_output_tokens += current_event_output_tokens
                        
                        logger.info(f"OpenAIRealtime Usage Parsed: Input Tokens Added: {current_event_input_tokens}, Output Tokens Added: {current_event_output_tokens}. Cumulative: Input={self.current_input_tokens}, Output={self.current_output_tokens}")
                        
                        self.raw_usage_events_for_turn.append({"type": event.type, "usage": dict(event.usage) if hasattr(event.usage, 'dict') else vars(event.usage)})

                    elif event.type == "error":
                        error_code = "N/A"
                        error_message_text = "Unknown error from OpenAI Realtime API."
                        if hasattr(event, 'error') and event.error:
                            if hasattr(event.error, 'code'):
                                error_code = event.error.code
                            if hasattr(event.error, 'message'):
                                error_message_text = event.error.message
                        else:
                            # Fallback if event.error structure is not as expected or missing
                            error_message_text = str(event) # Convert the event to string as a last resort

                        full_error_details = f"Code: {error_code}, Message: {error_message_text}"
                        logger.error(f"OpenAI Realtime API Error: {full_error_details}")
                        
                        await self.output_queue.put(
                            AdditionalOutputs(
                                {
                                    "type": "status_update",
                                    "status": "error",
                                    "message": f"OpenAI Error: {error_message_text}", # UI gets the message part
                                }
                            )
                        )
                        error_chat_message = ChatMessage(role="assistant", content=f"[OpenAI Error: {error_message_text}]")
                        await self.output_queue.put(
                            AdditionalOutputs({"type": "chatbot_update", "message": error_chat_message.model_dump()})
                        )
                        await self.output_queue.put(
                            AdditionalOutputs(
                                {
                                    "type": "status_update",
                                    "status": "idle",
                                    "message": "Ready (after error)",
                                }
                            )
                        )
                        # Removed problematic final_chatbot_state
        except openai.AuthenticationError as e:
            logger.error(f"OpenAIRealtimeHandler: Authentication Error: {e}. Check your --openai-api-key.")
            await self.output_queue.put(AdditionalOutputs({"type": "status_update", "status": "error", "message": "OpenAI Auth Error. Check API Key."}))
        except Exception as e:
            logger.error(f"OpenAIRealtimeHandler: Connection failed or error during session: {e}", exc_info=True)
            await self.output_queue.put(AdditionalOutputs({"type": "status_update", "status": "error", "message": f"Connection Error: {str(e)}"}))
        finally:
            logger.info("OpenAIRealtimeHandler: start_up processing loop finished.")
            if self.connection:
                logger.info("OpenAIRealtimeHandler: Closing connection in start_up finally block.")
                try:
                    await self.connection.close()
                except Exception as e:
                    logger.warning(f"OpenAIRealtimeHandler: Error closing connection in start_up finally: {e}")
            self.connection = None 
            await self.output_queue.put(None) 


    async def receive(self, frame: tuple[int, np.ndarray]) -> None:
        if not self.connection: 
            return
        
        _, array = frame
        if array.ndim > 1: 
            array = array.squeeze()

        if array.dtype != np.int16:
            array = array.astype(np.int16)

        audio_bytes = array.tobytes()
        audio_message = base64.b64encode(audio_bytes).decode("utf-8")
        try:
            await self.connection.input_audio_buffer.append(audio=audio_message)
        except Exception as e: 
            logger.error(f"OpenAIRealtimeHandler: Error sending audio: {e}")


    async def emit(self) -> tuple[int, np.ndarray] | AdditionalOutputs | None:
        return await wait_for_item(self.output_queue)

    async def shutdown(self) -> None:
        logger.info("OpenAIRealtimeHandler: Shutting down...")
        if self.connection:
            logger.info("OpenAIRealtimeHandler: Closing connection in shutdown.")
            try:
                await self.connection.close()
            except Exception as e:
                logger.warning(f"OpenAIRealtimeHandler: Error closing connection in shutdown: {e}")
        self.connection = None 
        if self.client:
            await self.client.close() 
            self.client = None
        self.clear_queue()
        await self.output_queue.put(None) 
        logger.info("OpenAIRealtimeHandler: Shutdown complete.")

    def clear_queue(self):
        while not self.output_queue.empty():
            try:
                self.output_queue.get_nowait()
            except asyncio.QueueEmpty:
                break
        logger.debug("OpenAIRealtimeHandler: Output queue cleared.")

# --- Endpoint Definitions ---
def register_endpoints(app: FastAPI, stream: Stream):
    """Registers FastAPI endpoints."""
    curr_dir = Path(__file__).parent

    @app.get("/")
    async def _():
        rtc_config = get_twilio_turn_credentials() if get_space() else None
        index_path = curr_dir / "index.html"
        if not index_path.exists():
            logger.error("index.html not found in the current directory!")
            return HTMLResponse(
                content="<html><body><h1>Error: index.html not found</h1></body></html>",
                status_code=500,
            )

        html_content = index_path.read_text()
        if rtc_config:
            html_content = html_content.replace(
                "__RTC_CONFIGURATION__", json.dumps(rtc_config)
            )
        else:
            html_content = html_content.replace("__RTC_CONFIGURATION__", "null")

        html_content = html_content.replace(
            "__SYSTEM_MESSAGE_JSON__", json.dumps(settings.system_message)
        )
        html_content = html_content.replace("__APP_VERSION__", APP_VERSION)
        html_content = html_content.replace(
            "__STT_LANGUAGE_JSON__", json.dumps(settings.current_stt_language)
        )
        # For TTS speed, OpenAI backend doesn't have a user-configurable speed via this app's settings.
        # It might be part of the voice model or a parameter not exposed here. Default to 1.0.
        tts_speed_to_inject = settings.current_tts_speed if settings.backend == "classic" else 1.0
        html_content = html_content.replace(
            "__TTS_SPEED_JSON__", json.dumps(tts_speed_to_inject)
        )
        html_content = html_content.replace(
            "__STARTUP_TIMESTAMP_STR__", json.dumps(settings.startup_timestamp_str)
        )
        html_content = html_content.replace(
            "__BACKEND_TYPE__", json.dumps(settings.backend)
        )
        return HTMLResponse(content=html_content, status_code=200)

    @app.post("/input_hook")
    async def _(body: InputData):
        chatbot_history = [msg.model_dump() for msg in body.chatbot]
        stream.set_input(body.webrtc_id, chatbot_history)
        return {"status": "ok"}

    @app.get("/outputs")
    def _(webrtc_id: str):
        async def output_stream():
            try:
                async for output in stream.output_stream(webrtc_id):
                    if isinstance(output, AdditionalOutputs):
                        data_payload = output.args[0]
                        if isinstance(data_payload, dict) and "type" in data_payload:
                            event_type = data_payload["type"]
                            try:
                                event_data_json = json.dumps(data_payload, ensure_ascii=False)
                                if settings.verbose:
                                    logger.debug(
                                        f"Sending SSE event: type={event_type}, data={event_data_json}..."
                                    )
                                else:
                                    logger.debug(
                                        f"Sending SSE event: type={event_type}, data={event_data_json[:100]}..."
                                    )
                                yield f"event: {event_type}\ndata: {event_data_json}\n\n"
                            except TypeError as e:
                                logger.error(
                                    f"Failed to serialize AdditionalOutputs payload to JSON: {e}. Payload: {data_payload}"
                                )
                        else:
                            logger.warning(
                                f"Received AdditionalOutputs with unexpected payload structure: {data_payload}"
                            )
                    elif (
                        isinstance(output, tuple)
                        and len(output) == 2
                        and isinstance(output[1], np.ndarray)
                    ):
                        logger.debug(
                            f"Output stream received audio tuple for webrtc_id {webrtc_id}, should be handled by track."
                        )
                        pass 
                    elif isinstance(output, bytes):
                        logger.warning(
                            "Received raw bytes directly in output stream, expected AdditionalOutputs or audio tuple via handler."
                        )
                    else:
                        logger.warning(
                            f"Received unexpected output type in stream: {type(output)}"
                        )
            except Exception as e:
                logger.error(f"Error in output stream for webrtc_id {webrtc_id}: {e}")
                try:
                    error_payload = {
                        "type": "error_event",
                        "message": f"Server stream error: {str(e)}",
                    }
                    yield f"event: error_event\ndata: {json.dumps(error_payload, ensure_ascii=False)}\n\n"
                except Exception as send_err:
                    logger.error(
                        f"Failed to send error event to client {webrtc_id}: {send_err}"
                    )
        return StreamingResponse(output_stream(), media_type="text/event-stream")

    @app.get("/available_models")
    async def get_available_models_endpoint(): 
        if settings.backend == "openai":
            return JSONResponse(
                {"available": [settings.openai_realtime_model_arg], "current": settings.openai_realtime_model_arg}
            )
        else: # classic
            return JSONResponse(
                {"available": settings.available_models, "current": settings.current_llm_model}
            )

    @app.get("/available_voices_tts")
    async def get_available_voices_tts_endpoint(): 
        if settings.backend == "openai":
            # OpenAI realtime backend uses its own set of voices, potentially configurable
            response_data = prepare_available_voices_data(
                settings.current_openai_voice, OPENAI_REALTIME_VOICES # Use OPENAI_REALTIME_VOICES
            )
            response_data["is_openai_realtime"] = True # Add a flag for UI
            return JSONResponse(response_data)
        else: # classic
            response_data = prepare_available_voices_data(
                settings.current_tts_voice, settings.available_voices_tts
            )
            response_data["is_openai_realtime"] = False
            return JSONResponse(response_data)

    @app.post("/switch_voice")
    async def switch_voice(request: Request):
        try:
            data = await request.json()
            new_voice_name = data.get("voice_name")
            if not new_voice_name:
                logger.warning("Missing voice_name in switch request.")
                return JSONResponse(
                    {"status": "error", "message": "Missing 'voice_name' in request body."},
                    status_code=400,
                )

            if settings.backend == "openai":
                if new_voice_name != settings.current_openai_voice:
                    if new_voice_name in OPENAI_REALTIME_VOICES:
                        settings.current_openai_voice = new_voice_name
                        logger.info(f"Switched active OpenAI realtime voice to: {settings.current_openai_voice}")
                        # The OpenAIRealtimeHandler will pick this up on the next connection.
                        return JSONResponse(
                            {"status": "success", "voice": settings.current_openai_voice}
                        )
                    else:
                        logger.warning(
                            f"Attempted to switch OpenAI realtime voice to '{new_voice_name}' which is not in the available list: {OPENAI_REALTIME_VOICES}"
                        )
                        return JSONResponse(
                            {"status": "error", "message": f"Voice '{new_voice_name}' is not available for OpenAI realtime backend."},
                            status_code=400,
                        )
                else:
                    logger.info(f"OpenAI realtime voice already set to: {new_voice_name}")
                    return JSONResponse(
                        {"status": "success", "voice": settings.current_openai_voice}
                    )
            else: # Classic backend logic
                if new_voice_name != settings.current_tts_voice:
                    if new_voice_name in settings.available_voices_tts:
                        settings.current_tts_voice = new_voice_name
                        logger.info(f"Switched active classic TTS voice to: {settings.current_tts_voice}")
                        return JSONResponse(
                            {"status": "success", "voice": settings.current_tts_voice}
                        )
                    else:
                        logger.warning(
                            f"Attempted to switch classic TTS voice to '{new_voice_name}' which is not in the available list: {settings.available_voices_tts}"
                        )
                        return JSONResponse(
                            {"status": "error", "message": f"Voice '{new_voice_name}' is not available for classic backend."},
                            status_code=400,
                        )
                else:
                    logger.info(f"Classic TTS voice already set to: {new_voice_name}")
                    return JSONResponse(
                        {"status": "success", "voice": settings.current_tts_voice}
                    ) 
        except json.JSONDecodeError:
            logger.error("Failed to decode JSON body in /switch_voice")
            return JSONResponse(
                {"status": "error", "message": "Invalid JSON format in request body"},
                status_code=400,
            )
        except Exception as e:
            logger.error(f"Error processing /switch_voice request: {e}")
            return JSONResponse(
                {"status": "error", "message": f"Internal server error: {str(e)}"},
                status_code=500,
            )

    @app.post("/switch_stt_language")
    async def switch_stt_language(request: Request):
        try:
            data = await request.json()
            new_language = data.get("stt_language", None) 

            if new_language is not None and not new_language.strip():
                new_language = None
            elif new_language is not None:
                new_language = new_language.strip()

            if new_language != settings.current_stt_language:
                settings.current_stt_language = new_language
                logger.info(
                    f"Switched active STT language to: '{settings.current_stt_language}' (None means auto-detect)"
                )
                # For OpenAI backend, the handler needs to pick up this change on next connection.
                return JSONResponse(
                    {"status": "success", "stt_language": settings.current_stt_language}
                )
            else:
                logger.info(
                    f"STT language already set to: '{settings.current_stt_language}'"
                )
                return JSONResponse(
                    {"status": "success", "stt_language": settings.current_stt_language}
                )
        except json.JSONDecodeError:
            logger.error("Failed to decode JSON body in /switch_stt_language")
            return JSONResponse(
                {"status": "error", "message": "Invalid JSON format in request body"},
                status_code=400,
            )
        except Exception as e:
            logger.error(f"Error processing /switch_stt_language request: {e}")
            return JSONResponse(
                {"status": "error", "message": f"Internal server error: {str(e)}"},
                status_code=500,
            )

    @app.post("/switch_tts_speed")
    async def switch_tts_speed(request: Request):
        if settings.backend == "openai":
            # OpenAI Realtime API might have speed control, but it's not exposed via this app's settings in the same way.
            # The 'output_audio_generation' parameters in OpenAI's API can include 'speed'.
            # For now, this app doesn't allow changing it for OpenAI backend.
            logger.info("TTS speed adjustment requested for OpenAI backend, which is not currently user-configurable via this app's UI.")
            return JSONResponse(
                {"status": "info", "message": "TTS speed adjustment for OpenAI realtime backend is handled by OpenAI or not user-configurable through this app."},
                status_code=200, # Not an error, just info.
            )
        # Classic backend logic
        try:
            data = await request.json()
            new_speed = data.get("tts_speed") 

            if new_speed is None:
                logger.warning("Missing 'tts_speed' in switch request.")
                return JSONResponse(
                    {"status": "error", "message": "Missing 'tts_speed' in request body."},
                    status_code=400,
                )
            try:
                new_speed_float = float(new_speed)
                if not (0.1 <= new_speed_float <= 4.0):
                    raise ValueError("TTS speed must be between 0.1 and 4.0")
            except (ValueError, TypeError):
                logger.warning(f"Invalid TTS speed value received: {new_speed}")
                return JSONResponse(
                    {"status": "error", "message": "Invalid TTS speed value. Must be a number between 0.1 and 4.0."},
                    status_code=400,
                )

            if new_speed_float != settings.current_tts_speed:
                settings.current_tts_speed = new_speed_float
                logger.info(
                    f"Switched active TTS speed (classic backend) to: {settings.current_tts_speed:.1f}"
                )
                return JSONResponse(
                    {"status": "success", "tts_speed": settings.current_tts_speed}
                )
            else:
                logger.info(
                    f"TTS speed (classic backend) already set to: {settings.current_tts_speed:.1f}"
                )
                return JSONResponse(
                    {"status": "success", "tts_speed": settings.current_tts_speed}
                )
        except json.JSONDecodeError:
            logger.error("Failed to decode JSON body in /switch_tts_speed")
            return JSONResponse(
                {"status": "error", "message": "Invalid JSON format in request body"},
                status_code=400,
            )
        except Exception as e:
            logger.error(f"Error processing /switch_tts_speed request: {e}")
            return JSONResponse(
                {"status": "error", "message": f"Internal server error: {str(e)}"},
                status_code=500,
            )

    @app.post("/switch_model")
    async def switch_model(request: Request):
        if settings.backend == "openai":
            # Model switching for OpenAI backend is not supported via this endpoint as it's set at startup.
            logger.warning("Attempt to switch model for OpenAI backend via API, which is not supported post-startup.")
            return JSONResponse(
                {"status": "error", "message": "Switching OpenAI realtime model via API is not supported. Model is fixed at startup."},
                status_code=400,
            )
        else: # classic backend
            try:
                data = await request.json()
                new_model_name = data.get("model_name")
                if new_model_name:
                    if new_model_name != settings.current_llm_model:
                        if new_model_name in settings.available_models:
                            settings.current_llm_model = new_model_name
                            logger.info(
                                f"Switched active LLM model (classic backend) to: {settings.current_llm_model}"
                            )
                            if (
                                new_model_name not in settings.model_cost_data
                                or settings.model_cost_data[new_model_name].get(
                                    "input_cost_per_token"
                                )
                                is None
                            ):
                                logger.warning(
                                    f"Cost data might be missing or incomplete for the newly selected model '{settings.current_llm_model}'."
                                )
                            return JSONResponse(
                                {"status": "success", "model": settings.current_llm_model}
                            )
                        else:
                            logger.warning(
                                f"Attempted to switch to model '{new_model_name}' which is not in the available list: {settings.available_models}"
                            )
                            return JSONResponse(
                                {"status": "error", "message": f"Model '{new_model_name}' is not available."},
                                status_code=400,
                            )
                    else:
                        logger.info(f"Model already set to: {new_model_name}")
                        return JSONResponse(
                            {"status": "success", "model": settings.current_llm_model}
                        )
                else: 
                    logger.warning(f"Missing model_name in switch request.")
                    return JSONResponse(
                        {"status": "error", "message": f"Missing 'model_name' in request body."},
                        status_code=400,
                    )
            except json.JSONDecodeError:
                logger.error("Failed to decode JSON body in /switch_model")
                return JSONResponse(
                    {"status": "error", "message": "Invalid JSON format in request body"},
                    status_code=400,
                )
            except Exception as e:
                logger.error(f"Error processing /switch_model request: {e}")
                return JSONResponse(
                    {"status": "error", "message": f"Internal server error: {str(e)}"},
                    status_code=500,
                )

    @app.post("/heartbeat")
    async def heartbeat(request: Request):
        global last_heartbeat_time
        try:
            last_heartbeat_time = datetime.datetime.now(datetime.timezone.utc)
            payload = await request.json()
            logger.debug(
                f"Heartbeat received at {last_heartbeat_time}. Payload: {payload}"
            )
            return {"status": "ok"}
        except json.JSONDecodeError:
            last_heartbeat_time = datetime.datetime.now(datetime.timezone.utc)
            logger.debug(
                f"Heartbeat received at {last_heartbeat_time} (no valid JSON payload)."
            )
            return {"status": "ok"}
        except Exception as e:
            logger.error(f"Error processing heartbeat: {e}")
            return JSONResponse({"status": "error", "message": str(e)}, status_code=500)

    @app.post("/reset_chat_log")
    async def reset_chat_log():
        try:
            new_timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            settings.startup_timestamp_str = new_timestamp
            logger.info(f"Chat log timestamp reset to: {settings.startup_timestamp_str}")
            # Also reset the TTS audio directory for the new "session"
            if settings.tts_base_dir: # Only relevant for classic backend
                settings.tts_audio_dir = settings.tts_base_dir / settings.startup_timestamp_str
                settings.tts_audio_dir.mkdir(exist_ok=True)
                logger.info(f"TTS audio directory for this session reset to: {settings.tts_audio_dir}")
            else:
                logger.warning("tts_base_dir not set, cannot reset TTS audio directory (or not applicable for current backend).")

            return JSONResponse(
                {"status": "success", "new_timestamp": settings.startup_timestamp_str}
            )
        except Exception as e:
            logger.error(f"Error resetting chat log timestamp: {e}")
            return JSONResponse(
                {"status": "error", "message": "Failed to reset chat log timestamp"},
                status_code=500,
            )

    @app.get("/tts_audio/{run_timestamp}/{filename}")
    async def get_tts_audio(run_timestamp: str, filename: str):
        if settings.backend == "openai":
            logger.warning(f"Request to /tts_audio for '{filename}' when using OpenAI backend. This endpoint is for classic backend.")
            raise HTTPException(status_code=404, detail="TTS audio file serving not applicable for OpenAI backend.")

        if not settings.tts_audio_dir: 
             logger.error("settings.tts_audio_dir not configured for classic backend, cannot serve audio.")
             raise HTTPException(status_code=500, detail="Server configuration error")

        if ".." in filename or "/" in filename or "\\" in filename:
            logger.warning(f"Attempt to access potentially unsafe filename: {filename}")
            raise HTTPException(status_code=400, detail="Invalid filename")
        
        file_path = settings.tts_audio_dir / filename
        logger.debug(f"Attempting to serve TTS audio file (classic backend): {file_path}")

        if file_path.is_file():
            return FileResponse(file_path, media_type="audio/mpeg", filename=filename)
        else:
            logger.warning(f"TTS audio file not found: {file_path}")
            raise HTTPException(status_code=404, detail="Audio file not found")

# --- Pywebview API Class ---
class Api:
    def __init__(self, window):
        self._window = window

    def close(self):
        logger.info("API close method called.")
        if self._window:
            self._window.destroy()

# --- Heartbeat Monitoring Thread ---
def monitor_heartbeat_thread():
    global last_heartbeat_time, uvicorn_server, pywebview_window, shutdown_event
    logger.info("Heartbeat monitor thread started.")
    initial_wait_done = False

    while not shutdown_event.is_set():
        if last_heartbeat_time is None:
            if not initial_wait_done:
                logger.info(
                    f"Waiting for the first heartbeat (timeout check in {heartbeat_timeout * 2}s)..."
                )
                shutdown_event.wait(heartbeat_timeout * 2)
                initial_wait_done = True
                if shutdown_event.is_set():
                    break 
                continue 
            else:
                logger.debug("Still waiting for first heartbeat...")
                shutdown_event.wait(5) 
                if shutdown_event.is_set():
                    break
                continue

        time_since_last = (
            datetime.datetime.now(datetime.timezone.utc) - last_heartbeat_time
        )
        logger.debug(
            f"Time since last heartbeat: {time_since_last.total_seconds():.1f}s"
        )

        if time_since_last.total_seconds() > heartbeat_timeout:
            logger.warning(
                f"Heartbeat timeout ({heartbeat_timeout}s exceeded). Initiating shutdown."
            )
            if uvicorn_server:
                logger.info("Signaling Uvicorn server to stop...")
                uvicorn_server.should_exit = True
            else:
                logger.warning(
                    "Uvicorn server instance not found, cannot signal shutdown."
                )

            if pywebview_window:
                logger.info("Destroying pywebview window...")
                try:
                    pywebview_window.destroy()
                except Exception as e:
                    logger.error(
                        f"Error destroying pywebview window from monitor thread: {e}"
                    )
            break
        shutdown_event.wait(5)
    logger.info("Heartbeat monitor thread finished.")


@click.command(help="Run a simple voice chat interface using a configurable LLM provider, STT server, and TTS.")
@click.option(
    "--host",
    type=str,
    default="127.0.0.1",
    show_default=True,
    help="Host address to bind the FastAPI server to.",
)
@click.option(
    "--port",
    type=click.INT,
    envvar="APP_PORT",
    default=int(APP_PORT_ENV), 
    show_default=True, 
    help="Preferred port to run the FastAPI server on. (Env: APP_PORT)",
)
@click.option(
    "-v",
    "--verbose",
    is_flag=True,
    default=False,
    show_default=True,
    help="Enable verbose logging (DEBUG level).",
)
@click.option(
    "--browser",
    is_flag=True,
    default=False,
    show_default=True,
    help="Launch the application in the default web browser instead of a dedicated GUI window.",
)
@click.option(
    "--system-message",
    type=str,
    envvar="SYSTEM_MESSAGE",
    default=SYSTEM_MESSAGE_ENV, 
    show_default=True,
    help="System message to prepend to the chat history. (Env: SYSTEM_MESSAGE)",
)
@click.option(
    "--backend",
    type=click.Choice(["classic", "openai"], case_sensitive=False),
    default="classic",
    show_default=True,
    help="Backend to use for voice processing. 'classic' uses separate STT/LLM/TTS. 'openai' uses OpenAI's realtime voice API.",
)
@click.option(
    "--openai-realtime-model",
    type=str,
    envvar="OPENAI_REALTIME_MODEL", 
    default=OPENAI_REALTIME_MODEL_ENV, 
    show_default=True,
    help="OpenAI realtime API model to use (if --backend=openai). (Env: OPENAI_REALTIME_MODEL)",
)
@click.option(
    "--openai-realtime-voice",
    type=str,
    envvar="OPENAI_REALTIME_VOICE",
    default=OPENAI_REALTIME_VOICE_ENV,
    show_default=True,
    help="Default voice for OpenAI realtime backend (if --backend=openai). (Env: OPENAI_REALTIME_VOICE)",
)
@click.option(
    "--openai-api-key",
    type=str,
    envvar="OPENAI_API_KEY",
    default=OPENAI_API_KEY_ENV,
    show_default=True,
    help="API key for OpenAI services (REQUIRED if --backend=openai). (Env: OPENAI_API_KEY)",
)
@click.option(
    "--llm-host",
    type=str,
    envvar="LLM_HOST",
    default=LLM_HOST_ENV, 
    show_default=True,
    help="Host address of the LLM proxy server (classic backend, optional). (Env: LLM_HOST)",
)
@click.option(
    "--llm-port",
    type=str, 
    envvar="LLM_PORT",
    default=LLM_PORT_ENV, 
    show_default=True,
    help="Port of the LLM proxy server (classic backend, optional). (Env: LLM_PORT)",
)
@click.option(
    "--llm-model",
    type=str,
    envvar="LLM_MODEL",
    default=DEFAULT_LLM_MODEL_ENV,
    show_default=True,
    help="Default LLM model to use for classic backend (e.g., 'gpt-4o', 'litellm_proxy/claude-3-opus'). (Env: LLM_MODEL)",
)
@click.option(
    "--llm-api-key",
    type=str,
    envvar="LLM_API_KEY",
    default=LLM_API_KEY_ENV, 
    show_default=True,
    help="API key for the LLM provider/proxy (classic backend, optional). (Env: LLM_API_KEY)",
)
@click.option(
    "--stt-host",
    type=str,
    envvar="STT_HOST",
    default=STT_HOST_ENV,
    show_default=True,
    help="Host address of the STT server (classic backend). (Env: STT_HOST)",
)
@click.option(
    "--stt-port",
    type=str, 
    envvar="STT_PORT",
    default=STT_PORT_ENV,
    show_default=True,
    help="Port of the STT server (classic backend). (Env: STT_PORT)",
)
@click.option(
    "--stt-model",
    type=str,
    envvar="STT_MODEL",
    default=STT_MODEL_ENV,
    show_default=True,
    help="STT model to use (classic backend). (Env: STT_MODEL)",
)
@click.option(
    "--stt-language",
    type=str,
    envvar="STT_LANGUAGE",
    default=STT_LANGUAGE_ENV, 
    show_default=True,
    help="Language code for STT (e.g., 'en', 'fr'). Used by both backends. If unset, Whisper usually auto-detects. (Env: STT_LANGUAGE)",
)
@click.option(
    "--stt-api-key",
    type=str,
    envvar="STT_API_KEY",
    default=STT_API_KEY_ENV, 
    show_default=True,
    help="API key for the STT server (classic backend, e.g., for OpenAI STT). (Env: STT_API_KEY)",
)
@click.option(
    "--stt-no-speech-prob-threshold",
    type=click.FLOAT,
    envvar="STT_NO_SPEECH_PROB_THRESHOLD",
    default=float(STT_NO_SPEECH_PROB_THRESHOLD_ENV),
    show_default=True,
    help=f"STT confidence (classic backend): Reject if no_speech_prob > this. (Env: STT_NO_SPEECH_PROB_THRESHOLD)",
)
@click.option(
    "--stt-avg-logprob-threshold",
    type=click.FLOAT,
    envvar="STT_AVG_LOGPROB_THRESHOLD",
    default=float(STT_AVG_LOGPROB_THRESHOLD_ENV),
    show_default=True,
    help=f"STT confidence (classic backend): Reject if avg_logprob < this. (Env: STT_AVG_LOGPROB_THRESHOLD)",
)
@click.option(
    "--stt-min-words-threshold",
    type=click.INT,
    envvar="STT_MIN_WORDS_THRESHOLD",
    default=int(STT_MIN_WORDS_THRESHOLD_ENV),
    show_default=True,
    help=f"STT confidence (classic backend): Reject if word count < this. (Env: STT_MIN_WORDS_THRESHOLD)",
)
@click.option(
    "--tts-host",
    type=str,
    envvar="TTS_HOST",
    default=TTS_HOST_ENV,
    show_default=True,
    help="Host address of the TTS server (classic backend). (Env: TTS_HOST)",
)
@click.option(
    "--tts-port",
    type=str, 
    envvar="TTS_PORT",
    default=TTS_PORT_ENV,
    show_default=True,
    help="Port of the TTS server (classic backend). (Env: TTS_PORT)",
)
@click.option(
    "--tts-model",
    type=str,
    envvar="TTS_MODEL",
    default=TTS_MODEL_ENV,
    show_default=True,
    help="TTS model to use (classic backend). (Env: TTS_MODEL)",
)
@click.option(
    "--tts-voice",
    type=str,
    envvar="TTS_VOICE", # This is for CLASSIC backend TTS voice
    default=DEFAULT_VOICE_TTS_ENV,
    show_default=True,
    help="Default TTS voice to use (classic backend). (Env: TTS_VOICE)",
)
@click.option(
    "--tts-api-key",
    type=str,
    envvar="TTS_API_KEY",
    default=TTS_API_KEY_ENV, 
    show_default=True,
    help="API key for the TTS server (classic backend, e.g., for OpenAI TTS). (Env: TTS_API_KEY)",
)
@click.option(
    "--tts-speed",
    type=click.FLOAT,
    envvar="TTS_SPEED",
    default=float(DEFAULT_TTS_SPEED_ENV),
    show_default=True,
    help=f"Default TTS speed multiplier (classic backend). (Env: TTS_SPEED)",
)
@click.option(
    "--tts-acronym-preserve-list",
    type=str,
    envvar="TTS_ACRONYM_PRESERVE_LIST",
    default=TTS_ACRONYM_PRESERVE_LIST_ENV, 
    show_default=True,
    help=f"Comma-separated list of acronyms to preserve during TTS (classic backend, Kokoro TTS). (Env: TTS_ACRONYM_PRESERVE_LIST)",
)
def main(
    host: str,
    port: int,
    verbose: bool,
    browser: bool,
    system_message: Optional[str],
    backend: str, 
    openai_realtime_model: str,
    openai_realtime_voice: str, # New CLI option for OpenAI backend voice
    openai_api_key: Optional[str], 
    llm_host: Optional[str],
    llm_port: Optional[str],
    llm_model: str,
    llm_api_key: Optional[str], 
    stt_host: str,
    stt_port: str,
    stt_model: str,
    stt_language: Optional[str], 
    stt_api_key: Optional[str], 
    stt_no_speech_prob_threshold: float,
    stt_avg_logprob_threshold: float,
    stt_min_words_threshold: int,
    tts_host: str,
    tts_port: str,
    tts_model: str,
    tts_voice: str, # This is for CLASSIC backend TTS voice
    tts_api_key: Optional[str], 
    tts_speed: float,
    tts_acronym_preserve_list: str,
) -> int:
    global uvicorn_server, pywebview_window

    startup_time = datetime.datetime.now()
    startup_timestamp_str_local = startup_time.strftime("%Y%m%d_%H%M%S")

    settings.startup_timestamp_str = startup_timestamp_str_local
    settings.backend = backend
    settings.openai_realtime_model_arg = openai_realtime_model
    settings.openai_realtime_voice_arg = openai_realtime_voice # Store initial arg
    settings.openai_api_key = openai_api_key 

    settings.preferred_port = port
    settings.host = host
    settings.verbose = verbose
    settings.browser = browser
    settings.system_message = system_message.strip() if system_message is not None else ""

    # --- Logging Setup (Early) ---
    console_log_level_str = "DEBUG" if settings.verbose else "INFO"
    log_file_path_for_setup: Optional[Path] = None
    log_dir_creation_error_details: Optional[str] = None 
    try:
        app_name = "SimpleVoiceChat"
        app_author = "Attila"
        log_base_dir_path_str: Optional[str] = None
        try:
            log_base_dir_path_str = platformdirs.user_log_dir(app_name, app_author)
        except Exception as e_log_dir:
            print(f"Warning: Could not find user log directory ({e_log_dir}), falling back to user data directory for logs.", file=sys.stderr)
            try:
                log_base_dir_path_str = platformdirs.user_data_dir(app_name, app_author)
            except Exception as e_data_dir:
                log_dir_creation_error_details = f"Could not find user data directory either ({e_data_dir})."
                print(f"Error: {log_dir_creation_error_details} File logging will be disabled.", file=sys.stderr)
        
        if log_base_dir_path_str:
            log_base_dir = Path(log_base_dir_path_str)
            settings.app_log_dir = log_base_dir / "logs"
            settings.app_log_dir.mkdir(parents=True, exist_ok=True)
            log_file_path_for_setup = settings.app_log_dir / f"log_{settings.startup_timestamp_str}.log"
        elif not log_dir_creation_error_details: 
            log_dir_creation_error_details = "Failed to determine a valid base directory for logs."
            print(f"Error: {log_dir_creation_error_details} File logging will be disabled.", file=sys.stderr)
            
    except Exception as e:
        log_dir_creation_error_details = f"Failed to set up log directory structure: {e}."
        print(f"Error: {log_dir_creation_error_details} File logging will be disabled.", file=sys.stderr)

    setup_logging(console_log_level_str, log_file_path_for_setup, settings.verbose)
    if log_dir_creation_error_details:
        logger.error(f"Log directory setup failed: {log_dir_creation_error_details} File logging is disabled.")
    
    logger.info(f"Application Version: {APP_VERSION}")
    logger.info(f"Using backend: {settings.backend}")

    # --- STT Language (Common to both backends) ---
    settings.stt_language_arg = stt_language 
    if settings.stt_language_arg:
        logger.info(f"STT language specified: {settings.stt_language_arg}")
        settings.current_stt_language = settings.stt_language_arg
    else:
        logger.info("No STT language specified (or empty), STT will auto-detect.")
        settings.current_stt_language = None

    # --- Backend Specific Configuration ---
    if settings.backend == "openai":
        logger.info(f"Configuring for 'openai' backend.")
        
        settings.available_models = OPENAI_REALTIME_MODELS # Use the list from config

        if not settings.available_models:
            logger.critical(
                "OPENAI_REALTIME_MODELS list is empty in configuration. "
                "Cannot proceed with OpenAI backend. Exiting."
            )
            return 1

        # Validate the chosen model (from CLI/env, stored in settings.openai_realtime_model_arg)
        # against the available list.
        chosen_model = settings.openai_realtime_model_arg
        if chosen_model in settings.available_models:
            settings.current_llm_model = chosen_model
        else:
            logger.warning(
                f"OpenAI realtime model '{chosen_model}' (from --openai-realtime-model or env) "
                f"is not in the list of supported models: {settings.available_models}. "
                f"Defaulting to the first available model: '{settings.available_models[0]}'."
            )
            settings.current_llm_model = settings.available_models[0]
        
        logger.info(f"OpenAI Realtime Model set to: {settings.current_llm_model}")
        settings.model_cost_data = {} # Cost for OpenAI realtime is handled differently (per token via OPENAI_REALTIME_PRICING)

        if not settings.openai_api_key:
            logger.critical("OpenAI API Key (--openai-api-key or OPENAI_API_KEY env) is REQUIRED for 'openai' backend. Exiting.")
            return 1
        logger.info("Using dedicated OpenAI API Key for 'openai' backend.")

        # OpenAI Realtime Voice
        settings.available_voices_tts = OPENAI_REALTIME_VOICES # Populate for UI dropdown
        initial_openai_voice_preference = settings.openai_realtime_voice_arg
        if initial_openai_voice_preference and initial_openai_voice_preference in OPENAI_REALTIME_VOICES:
            settings.current_openai_voice = initial_openai_voice_preference
        elif OPENAI_REALTIME_VOICES:
            if initial_openai_voice_preference:
                 logger.warning(f"OpenAI realtime voice '{initial_openai_voice_preference}' not found. Using first available: {OPENAI_REALTIME_VOICES[0]}.")
            settings.current_openai_voice = OPENAI_REALTIME_VOICES[0]
        else: # Should not happen if OPENAI_REALTIME_VOICES is populated
            settings.current_openai_voice = initial_openai_voice_preference 
            logger.error(f"No OpenAI realtime voices available or specified voice '{settings.current_openai_voice}' is invalid. Voice may not work.")
        logger.info(f"Initial OpenAI realtime voice set to: {settings.current_openai_voice}")


        # Log warnings if classic backend STT/TTS params are provided unnecessarily
        if stt_host != STT_HOST_ENV or stt_port != STT_PORT_ENV or stt_model != STT_MODEL_ENV or stt_api_key is not None:
            logger.warning("STT host/port/model/api-key parameters are ignored when using 'openai' backend (except STT language).")
        if tts_host != TTS_HOST_ENV or tts_port != TTS_PORT_ENV or tts_model != TTS_MODEL_ENV or tts_voice != DEFAULT_VOICE_TTS_ENV or tts_api_key is not None or tts_speed != float(DEFAULT_TTS_SPEED_ENV):
            logger.warning("Classic TTS host/port/model/voice/api-key/speed parameters are ignored when using 'openai' backend.")
        
        settings.current_tts_speed = 1.0 # Not user-configurable for OpenAI backend via this app

    elif settings.backend == "classic":
        logger.info(f"Configuring for 'classic' backend.")
        # --- LLM Configuration (Classic) ---
        settings.llm_host_arg = llm_host
        settings.llm_port_arg = llm_port
        settings.llm_model_arg = llm_model 
        settings.llm_api_key = llm_api_key 

        settings.use_llm_proxy = bool(settings.llm_host_arg and settings.llm_port_arg)
        if settings.use_llm_proxy:
            try:
                llm_port_int = int(settings.llm_port_arg) # type: ignore
                settings.llm_api_base = f"http://{settings.llm_host_arg}:{llm_port_int}/v1"
                logger.info(f"Using LLM proxy at: {settings.llm_api_base}")
                if settings.llm_api_key:
                    logger.info("Using LLM API key for proxy.")
                else:
                    logger.info("No LLM API key provided for proxy (assumed optional).")
            except (ValueError, TypeError):
                logger.error(
                    f"Error: Invalid LLM port specified: '{settings.llm_port_arg}'. Disabling proxy."
                )
                settings.use_llm_proxy = False
                settings.llm_api_base = None
        else:
            settings.llm_api_base = None
            logger.info("Not using LLM proxy (using default LLM routing).")
            if settings.llm_api_key:
                logger.info("Using LLM API key for direct routing.")
            else:
                logger.info("No LLM API key provided for direct routing (will use LiteLLM's environment config).")

        # --- STT Configuration (Classic) ---
        settings.stt_host_arg = stt_host
        settings.stt_port_arg = stt_port
        settings.stt_model_arg = stt_model
        settings.stt_api_key = stt_api_key 
        settings.stt_no_speech_prob_threshold = stt_no_speech_prob_threshold
        settings.stt_avg_logprob_threshold = stt_avg_logprob_threshold
        settings.stt_min_words_threshold = stt_min_words_threshold

        settings.is_openai_stt = settings.stt_host_arg == "api.openai.com"
        if settings.is_openai_stt:
            settings.stt_api_base = "https://api.openai.com/v1"
            logger.info(f"Using OpenAI STT at: {settings.stt_api_base} with model {settings.stt_model_arg}")
            if not settings.stt_api_key:
                logger.critical(
                    "STT_API_KEY (--stt-api-key) is required when using OpenAI STT (stt-host=api.openai.com) with classic backend. Exiting."
                )
                return 1
            logger.info("Using STT API key for OpenAI STT.")
        else: 
            try:
                stt_port_int = int(settings.stt_port_arg)
                scheme = "http" 
                settings.stt_api_base = f"{scheme}://{settings.stt_host_arg}:{stt_port_int}/v1"
                logger.info(f"Using Custom STT server at: {settings.stt_api_base} with model {settings.stt_model_arg}")
                if settings.stt_api_key:
                    logger.info("Using STT API key for custom STT server.")
                else:
                    logger.info("No STT API key provided for custom STT server (assumed optional).")
            except (ValueError, TypeError):
                logger.critical(
                    f"Invalid STT port specified for custom server: '{settings.stt_port_arg}'. Cannot connect. Exiting."
                )
                return 1
        logger.info(
            f"STT Confidence Thresholds: no_speech_prob > {settings.stt_no_speech_prob_threshold}, avg_logprob < {settings.stt_avg_logprob_threshold}, min_words < {settings.stt_min_words_threshold}"
        )

        # --- TTS Configuration (Classic) ---
        settings.tts_host_arg = tts_host
        settings.tts_port_arg = tts_port
        settings.tts_model_arg = tts_model
        settings.tts_voice_arg = tts_voice # Classic backend TTS voice
        settings.tts_api_key = tts_api_key 
        settings.tts_speed_arg = tts_speed 
        settings.tts_acronym_preserve_list_arg = tts_acronym_preserve_list

        settings.is_openai_tts = settings.tts_host_arg == "api.openai.com"
        if settings.is_openai_tts:
            settings.tts_base_url = "https://api.openai.com/v1"
            logger.info(f"Using OpenAI TTS at: {settings.tts_base_url} with model {settings.tts_model_arg}")
            if not settings.tts_api_key:
                logger.critical(
                    "TTS_API_KEY (--tts-api-key) is required when using OpenAI TTS (tts-host=api.openai.com) with classic backend. Exiting."
                )
                return 1
            logger.info("Using TTS API key for OpenAI TTS.")
            if settings.tts_model_arg in OPENAI_TTS_PRICING:
                logger.info(
                    f"OpenAI TTS pricing for '{settings.tts_model_arg}': ${OPENAI_TTS_PRICING[settings.tts_model_arg]:.2f} / 1M chars"
                )
        else: 
            try:
                tts_port_int = int(settings.tts_port_arg)
                scheme = "http" 
                settings.tts_base_url = f"{scheme}://{settings.tts_host_arg}:{tts_port_int}/v1"
                logger.info(f"Using Custom TTS server at: {settings.tts_base_url} with model {settings.tts_model_arg}")
                if settings.tts_api_key:
                    logger.info("Using TTS API key for custom TTS server.")
                else:
                    logger.info("No TTS API key provided for custom TTS server (assumed optional).")
            except (ValueError, TypeError):
                logger.critical(
                    f"Invalid TTS port specified for custom server: '{settings.tts_port_arg}'. Cannot connect. Exiting."
                )
                return 1
        
        settings.tts_acronym_preserve_set = {
            word.strip().upper()
            for word in settings.tts_acronym_preserve_list_arg.split(",")
            if word.strip()
        }
        logger.debug(f"Loaded TTS_ACRONYM_PRESERVE_SET: {settings.tts_acronym_preserve_set}")
        settings.current_tts_speed = settings.tts_speed_arg
        logger.info(f"Initial TTS speed (classic backend): {settings.current_tts_speed:.1f}")


        # --- Initialize Clients (Classic Backend) ---
        logger.info("Initializing clients for 'classic' backend...")
        try:
            settings.stt_client = OpenAI(
                base_url=settings.stt_api_base,
                api_key=settings.stt_api_key, 
            )
            logger.info(f"STT client initialized for classic backend (target: {settings.stt_api_base}).")
        except Exception as e:
            logger.critical(f"Failed to initialize STT client for classic backend: {e}. Exiting.")
            return 1
        
        try:
            settings.tts_client = OpenAI(
                base_url=settings.tts_base_url,
                api_key=settings.tts_api_key, 
            )
            logger.info(f"TTS client initialized for classic backend (target: {settings.tts_base_url}).")
        except Exception as e:
            logger.critical(f"Failed to initialize TTS client for classic backend: {e}. Exiting.")
            return 1

        # --- Model & Voice Availability (Classic Backend) ---
        if settings.use_llm_proxy:
            settings.available_models, settings.model_cost_data = get_models_and_costs_from_proxy(
                settings.llm_api_base, settings.llm_api_key 
            )
        else:
            settings.available_models, settings.model_cost_data = get_models_and_costs_from_litellm()

        if not settings.available_models:
            logger.warning("No LLM models found from proxy or litellm. Using fallback.")
            settings.available_models = ["fallback/unknown-model"]

        initial_llm_model_preference = settings.llm_model_arg
        if initial_llm_model_preference and initial_llm_model_preference in settings.available_models:
            settings.current_llm_model = initial_llm_model_preference
        elif settings.available_models and settings.available_models[0] != "fallback/unknown-model":
            if initial_llm_model_preference:
                logger.warning(f"LLM model '{initial_llm_model_preference}' not found. Using first available: {settings.available_models[0]}.")
            settings.current_llm_model = settings.available_models[0]
        else:
            settings.current_llm_model = initial_llm_model_preference or "fallback/unknown-model"
            logger.warning(f"Using specified or fallback LLM model: {settings.current_llm_model}. Availability/cost data might be missing.")
        logger.info(f"Initial LLM model (classic backend) set to: {settings.current_llm_model}")
        
        if settings.is_openai_tts:
            settings.available_voices_tts = OPENAI_TTS_VOICES
        else: 
            settings.available_voices_tts = get_voices(settings.tts_base_url, settings.tts_api_key)
            if not settings.available_voices_tts:
                logger.warning(f"Could not retrieve voices from custom TTS server at {settings.tts_base_url}.")
        logger.info(f"Available TTS voices (classic backend): {settings.available_voices_tts}")

        initial_tts_voice_preference = settings.tts_voice_arg # Classic TTS voice
        if initial_tts_voice_preference and initial_tts_voice_preference in settings.available_voices_tts:
            settings.current_tts_voice = initial_tts_voice_preference
        elif settings.available_voices_tts:
            if initial_tts_voice_preference:
                 logger.warning(f"Classic TTS voice '{initial_tts_voice_preference}' not found. Using first available: {settings.available_voices_tts[0]}.")
            settings.current_tts_voice = settings.available_voices_tts[0]
        else:
            settings.current_tts_voice = initial_tts_voice_preference
            logger.error(f"No classic TTS voices available or specified voice '{settings.current_tts_voice}' is invalid. TTS may fail.")
        logger.info(f"Initial classic TTS voice set to: {settings.current_tts_voice}")

    # --- Common Post-Backend-Specific Setup ---
    if settings.system_message:
        logger.info(f"Loaded SYSTEM_MESSAGE: '{settings.system_message[:50]}...'")
    else:
        logger.info("No SYSTEM_MESSAGE defined.")

    try:
        app_name = "SimpleVoiceChat"
        app_author = "Attila"
        user_data_dir = Path(platformdirs.user_data_dir(app_name, app_author))
        settings.chat_log_dir = user_data_dir / "chats"
        settings.chat_log_dir.mkdir(parents=True, exist_ok=True)
        logger.debug(f"Chat log directory set to: {settings.chat_log_dir}")
    except Exception as e:
        logger.error(f"Failed to create chat log directory: {e}. Chat logging disabled.")
        settings.chat_log_dir = None

    # TTS audio directory setup (primarily for classic backend)
    if settings.backend == "classic":
        try:
            app_name = "SimpleVoiceChat"
            app_author = "Attila"
            try:
                cache_base_dir = Path(platformdirs.user_cache_dir(app_name, app_author))
            except Exception:
                logger.warning("Could not find user cache directory, falling back to user data directory for TTS audio.")
                cache_base_dir = Path(platformdirs.user_data_dir(app_name, app_author))

            settings.tts_base_dir = cache_base_dir / "tts_audio"
            settings.tts_base_dir.mkdir(parents=True, exist_ok=True)
            logger.debug(f"Base TTS audio directory: {settings.tts_base_dir}")

            settings.tts_audio_dir = settings.tts_base_dir / settings.startup_timestamp_str
            settings.tts_audio_dir.mkdir(exist_ok=True)
            logger.info(f"This run's TTS audio directory (classic backend): {settings.tts_audio_dir}")
        except Exception as e:
            logger.error(f"Failed to create temporary TTS audio directory for classic backend: {e}. TTS audio saving might fail.")
    else: 
        logger.info("TTS audio file saving to disk is not applicable for 'openai' backend.")
        settings.tts_audio_dir = None 


    logger.info(f"Application server host: {settings.host}")
    logger.info(
        f"Application server preferred port: {settings.preferred_port}"
    )

    # --- Stream Handler Setup ---
    stream_handler: Any 
    if settings.backend == "openai":
        logger.info("Initializing Stream with OpenAIRealtimeHandler.")
        stream_handler = OpenAIRealtimeHandler(app_settings=settings)
    else: 
        logger.info("Initializing Stream with ReplyOnPause handler for classic backend.")
        stream_handler = ReplyOnPause(
            response, 
            algo_options=AlgoOptions(
                audio_chunk_duration=3.0,
                started_talking_threshold=0.2,
                speech_threshold=0.2,
            ),
            model_options=SileroVadOptions(
                threshold=0.6,
                min_speech_duration_ms=800,
                min_silence_duration_ms=3500,
            ),
            can_interrupt=True,
        )

    stream = Stream(
        modality="audio",
        mode="send-receive",
        handler=stream_handler,
        track_constraints={
            "echoCancellation": True,
            "noiseSuppression": {"exact": True}, 
            "autoGainControl": {"exact": True},  
            "sampleRate": {"ideal": OPENAI_REALTIME_SAMPLE_RATE}, 
            "sampleSize": {"ideal": 16},
            "channelCount": {"exact": 1},
        },
        rtc_configuration=get_twilio_turn_credentials() if get_space() else None,
        concurrency_limit=5 if get_space() else None,
        time_limit=180 if get_space() else None, 
    )

    app = FastAPI()
    stream.mount(app)
    register_endpoints(app, stream)

    current_host = settings.host 
    preferred_port_val = settings.preferred_port
    actual_port = preferred_port_val
    max_retries = 10

    if is_port_in_use(actual_port, current_host):
        logger.warning(
            f"Preferred port {actual_port} on host {current_host} is in use. Searching for an available port..."
        )
        found_port = False
        for attempt in range(max_retries):
            new_port = random.randint(1024, 65535)
            logger.debug(f"Attempt {attempt+1}: Checking port {new_port} on {current_host}...")
            if not is_port_in_use(new_port, current_host):
                actual_port = new_port
                found_port = True
                logger.info(f"Found available port: {actual_port} on host {current_host}")
                break
        if not found_port:
            logger.error(
                f"Could not find an available port on host {current_host} after {max_retries} attempts. Exiting."
            )
            return 1
    else:
        logger.info(f"Using preferred port {actual_port} on host {current_host}")

    settings.port = actual_port 
    url = f"http://{current_host}:{actual_port}"

    def run_server():
        global uvicorn_server
        try:
            config = uvicorn.Config(
                app,
                host=current_host, 
                port=actual_port,   
                log_config=None,
            )
            uvicorn_server = uvicorn.Server(config)
            logger.info(f"Starting Uvicorn server on {current_host}:{actual_port}...")
            uvicorn_server.run()
            logger.info("Uvicorn server has stopped.")
        except Exception as e:
            logger.critical(f"Uvicorn server encountered an error: {e}")
        finally:
            uvicorn_server = None

    monitor_thread = threading.Thread(target=monitor_heartbeat_thread, daemon=True)
    monitor_thread.start()

    server_thread = threading.Thread(target=run_server, daemon=True)
    server_thread.start()

    logger.debug("Waiting for Uvicorn server to initialize...")
    time.sleep(3.0)

    if not server_thread.is_alive() or uvicorn_server is None:
        logger.critical(
            "Uvicorn server thread failed to start or initialize correctly. Exiting."
        )
        return 1
    else:
        logger.debug("Server thread appears to be running.")

    exit_code = 0
    try:
        if settings.browser:
            logger.info(f"Opening application in default web browser at: {url}")
            webbrowser.open(url, new=1)
            logger.info(
                "Application opened in browser. Server is running in the background."
            )
            logger.info("Press Ctrl+C to stop the server.")
            try:
                server_thread.join()
            except KeyboardInterrupt:
                logger.info("KeyboardInterrupt received, shutting down.")
            finally:
                logger.info("Signaling heartbeat monitor thread to stop...")
                shutdown_event.set()
                if (
                    uvicorn_server
                    and server_thread.is_alive()
                    and not uvicorn_server.should_exit
                ):
                    logger.info("Signaling Uvicorn server to shut down...")
                    uvicorn_server.should_exit = True
                elif uvicorn_server and uvicorn_server.should_exit:
                    logger.info("Uvicorn server already signaled to shut down.")
                elif not server_thread.is_alive():
                    logger.info("Server thread already stopped.")
                else:
                    logger.warning(
                        "Uvicorn server instance not found, cannot signal shutdown."
                    )
                if server_thread.is_alive():
                    logger.info("Waiting for Uvicorn server thread to join...")
                    server_thread.join(timeout=5.0)
                    if server_thread.is_alive():
                        logger.warning(
                            "Uvicorn server thread did not exit gracefully after 5 seconds."
                        )
                    else:
                        logger.info("Uvicorn server thread joined successfully.")
                logger.info("Waiting for heartbeat monitor thread to join...")
                monitor_thread.join(timeout=2.0)
                if monitor_thread.is_alive():
                    logger.warning(
                        "Heartbeat monitor thread did not exit gracefully after 2 seconds."
                    )
                else:
                    logger.info("Heartbeat monitor thread joined successfully.")
        else:
            logger.info(f"Creating pywebview window for URL: {url}")
            api = Api(None)
            webview.settings['OPEN_DEVTOOLS_IN_DEBUG'] = False
            logger.info(f"pywebview setting OPEN_DEVTOOLS_IN_DEBUG set to False.")
            
            pywebview_window = webview.create_window(
                f"Simple Voice Chat v{APP_VERSION}",
                url,
                width=800,
                height=800,
                js_api=api,
            )
            api._window = pywebview_window

            logger.info("Starting pywebview...")
            try:
                webview.start(debug=True, gui='qt')
            except Exception as e:
                logger.critical(f"Pywebview encountered an error: {e}")
                exit_code = 1
            finally:
                logger.info("Pywebview window closed or heartbeat timed out.")
                logger.info("Signaling heartbeat monitor thread to stop...")
                shutdown_event.set()
                if uvicorn_server and not uvicorn_server.should_exit:
                    logger.info("Signaling Uvicorn server to shut down...")
                    uvicorn_server.should_exit = True
                elif uvicorn_server and uvicorn_server.should_exit:
                    logger.info("Uvicorn server already signaled to shut down.")
                else:
                    logger.warning(
                        "Uvicorn server instance not found, cannot signal shutdown."
                    )
                logger.info("Waiting for Uvicorn server thread to join...")
                server_thread.join(timeout=5.0)
                if server_thread.is_alive():
                    logger.warning(
                        "Uvicorn server thread did not exit gracefully after 5 seconds."
                    )
                else:
                    logger.info("Uvicorn server thread joined successfully.")
                logger.info("Waiting for heartbeat monitor thread to join...")
                monitor_thread.join(timeout=2.0)
                if monitor_thread.is_alive():
                    logger.warning(
                        "Heartbeat monitor thread did not exit gracefully after 2 seconds."
                    )
                else:
                    logger.info("Heartbeat monitor thread joined successfully.")
    except Exception as e:
        logger.critical(
            f"An unexpected error occurred in the main execution block: {e}",
            exc_info=True,
        )
        exit_code = 1
        shutdown_event.set()
        if uvicorn_server and not uvicorn_server.should_exit:
            uvicorn_server.should_exit = True

    logger.info(f"Main function returning exit code: {exit_code}")
    return exit_code

