ai-cli

#!/usr/bin/env python3
# PYTHON_ARGCOMPLETE_OK
import argparse
import argcomplete
import json
import miniaudio
import os
import psutil
import re
import sys
import tempfile
import time
import threading
import subprocess
import queue
import warnings
import utils
from utils import print_verbose
from utils import print_error
from byte_queue_file import ByteQueueFile
from collections import OrderedDict
from ai_provider import get_ai_providers
from tts_provider import get_tts_providers, TTSProvider
from stt_provider import get_stt_providers
from tts_optimizer import *
from version import *

warnings.simplefilter("ignore", UserWarning)

app_identifier = 'ai-cli'

models = {}
ai_providers = []

tts_models = {}
tts_providers = []
stt_providers = []
default_stt_provider = None

MIN_SENTENCE_LENGTH = 50
MAX_SENTENCE_LENGTH = 1000

LAST_PLAYED_TIMESTAMP = 0
AUDIO_TIMEOUT_SECONDS = 10

SILENCE_500_BYTES = b'ID3\x04\x00\x00\x00\x00\x00#TSSE\x00\x00\x00\x0f\x00\x00\x03Lavf60.16.100\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xff\xfbT\xc0\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00Info\x00\x00\x00\x0f\x00\x00\x00\x16\x00\x00\t\x00\x00    *****5555@@@@@JJJJUUUUU````jjjjjuuuu\x80\x80\x80\x80\x80\x8a\x8a\x8a\x8a\x95\x95\x95\x95\x95\xa0\xa0\xa0\xa0\xa0\xaa\xaa\xaa\xaa\xb5\xb5\xb5\xb5\xb5\xc0\xc0\xc0\xc0\xca\xca\xca\xca\xca\xd5\xd5\xd5\xd5\xe0\xe0\xe0\xe0\xe0\xea\xea\xea\xea\xf5\xf5\xf5\xf5\xf5\xff\xff\xff\xff\x00\x00\x00\x00Lavc60.31\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00$\x03\x00\x00\x00\x00\x00\x00\x00\t\x00\xa6\x83\xcf\xcb\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xff\xfb\x14\xc4\x00\x00C\xe0\x00\xf6\x00\x8c`\x00v\x00!p\x10\x890t\x1f|?\x82\x0b>Mg\xce\x07\xe2w\x14\x04\x01\x07\x14\xf8\x80\x10\x0c\x90\xe7\xeat\x10\xe0\x80 \x18L\x08P$\xb0\x85Z\xe61#\xd5\\e\ru\xc1 \xa2\x86\x0b\\\xbf\xa9\xfa\x93\xd0i\xcd\xbe\xb5\xd1\xd4\xdam\xd5\xd0\x01a\xee:\xb4\x16s\xd2\xff\xfb\x14\xc4\x07\x00\x03\x8c\x01\x15\x80\x84K\x80r\x00!l\x10\x88\x00\x81\xf2v\\\xabl_\xf6\x9e\xf7\xb3OH\x95\x02B\x10$\x84\x01\x83G\x10{\x90\xbb\xecS\x99d\xe6Y\x9fo\xb48\xea\xbbT\xa9\x02\xee,\x94\x04\x03\x9c@\xb2\x1c\\V\xf1\x8btQ\x8fB(D\xe1\x91|\xaa\xf6!\x88\x86\xc9\xbedl&\xb0\x93\xff\xfb\x14\xc4\x10\x00D$\x01\x08\xa0\x88I\x80j\x80a\x14\x00\x88\x00\x0cA(\n\x8b\xc6C\x17\xdfa6\x00\x1a\xd6\xda\x8b|\xf6\xea\x15\xae\xf1\xe8M\xeaj\xc7&\r\xb3kM-$\x04Rt&\x81:h\xaa2\x9e\xde\xe4%?\xca!\x9f\xde\xaa\x96\x08\t\x00\x00#j\x18\x9c\xcb\xca\xc8\xad=j|y\xc5\x92u{\xa3\xff\xfb\x14\xc4\x17\x80\x03<\x03\x11\x80\x04`\x00\x7f\x80! \x10\x8dp\xad\x13\xe120\xa9\x01\x9cP0>\xa4\x14V\xc0L\x15\xb1\x8f|\xd2\x96xYTO?-gE\xbcb,]\x8f\xa9[\\\x94HP\x02F(\xb7<\xa0\xc6"\xc7\xb9\rO\x7f\x1d\x89(\xb4\xae\x8e\xf3$%\xee\xb1\x95\x13:a\x99`2\xa5\x16A\xff\xfb\x14\xc4 \x00Cp\x03\r\x00\x84`\x00r\x00!\xb0\x11\x8dpw\xadDg\xc1^MU\xefmL\x9f\xcd\xec\xa7\xf0"\x9bL)\x08\x10K\x10\x84\x10`\x94pM\x1c\xf3\xe7\\\xb7{\xe5\x12\xab\xaf\x1d\x10.\xc2\xcbb\x89\xa5-\x13\xd5\xa1\x14bd$\xac\x03\xcd-K(\x95\xb5R5)\xfd\x95\x1f\xa9T\xed\xec:\xff\xfb\x14\xc4)\x80CT\x03\r\x00\x8c`\x00~\x00! \x10\x8d0\xdf@x,\x97V\x99)\x91\x11\x1a\x00\x01"U\x16j/B\x11\xb7e\x15\xe9\xd8\x98c\xee\xf6\xee\xa2\x85\x99\xd5\x96\xa8y\xb1t.\x99\x07\xa9\xf6)~\xb5Y\xbf\x96\xcbk\xf4\xa0Z\x8bE\x14\x80D\x00\x1a9`0\x83\x0f\xd2\x95,\x9a\xe3\xf2\xbc,\xff\xfb\x14\xc41\x80\x03\xbc\x03\x0f\x80\x04@\x00f\x80a\xf0\x00\x8c\x00\x9duaF\xb5\xa6Y\xb2}\xec{\x14\n\x9b@\x08\x00\x00)\xa3\xd2\x98\xd0f\r\xc8\xef\x97< \x91V^s\xbb\x94\xf2\xd5\x9a\xff\xf7\x88\x9a)sPl\xcbn\xa9>\x9c\xd1\xbc\xeb\x89\xd8\x1d^\xfe\n\x7f\xe2\x9b\xae\xeb\xf8\xde\xbf\xfe\xb7C\xdb\xa1y\xff\xfb\x14\xc4;\x00\x03(\x03\r\x00\x00\x00\x00\x7f\x00! \x10\x890\xe2\r\xe4;*!l\xba\xdb\xa7!\x00\x00\x00k\x960@\x12u\xe9\x8f\x91\xc9\xb4Wt\xb2P\x92\xec*\x08\xe8cZ\x95p\xeczC(\x07\x08\x05\xcbLJ\x08\x02\xb3\x82\xcf\x17\x9bI\xbbJ4\xc7eD\xab<cI\xa7\xc5\xfa\rx\xaa<\xb2 \xff\xfb\x14\xc4D\x00\xc4\x94\xb3\x05\x00\x84m\xc8\x89\x80 \x80\x00\x8c\x01#\x90+0\xc4Z\tAT\xcb\xcd\xed\xd7$\xec\xf9\xf4n-\xb3h\xed\x16\xb7\xe70I\x7f\xadnpd\x06\x8a\x08\xe3\xce.\x15b\x16\xba25\xb5\x8eE,\x14\x06\x07(Li\x07\x95\xa4\x9f\xaa\xb1\xc8r\xe4\xea\x80\x00\x01\xdd\xa9\x02\xaf\xcdo[^\xff\xfb\x14\xc4F\x00\x04D\x01\x07\x00\x84G\x00v\x00a\xa4\x00\x8c\x00\x95\xeb\x9c\xd9\xcf\xdeR#z\xcaj\xa6\\9\xe4\xe3O<eq\x01\x00\xb6\xf5e\xb3\xc5g\xc7\xfd\xb2?\xfe\xcf\xce\x9a\xb6\xd6]\xff\xe3\x7f\xe6\xf5\x8ft\x88\xe6x\xcc\xd9*\xbb\xb5Wj\t\x00B\x00$\xc5M\xd1R\xf2\xcbF\x8a\x9fzyy\xe7B\xff\xfb\x14\xc4K\x80D$\xc3\x06 \x84o\xc0|\x80! \x10\x8c\x00d\x14m\xf6\xa1\xc4\xc4L\xb8r\x03j\x00 ;\x96\x91C\xa1\x14\xf2\xfa\xe9r\xbf\x1b\xb2\x16\xf2\xac\x9b\xef5jlMu\xddu\xa3\x18\x12\xd5h\x04\x00\x00[\x86\xbd\n\x84\x16.,\x07[W}Ds\xcb\xb9~\xfb\xf9\x04\xdc\xe5G\x10\x94^\x10*\xff\xfb\x14\xc4P\x80\xc3\xc0\xc7\x08\xa0\x84Q\xd0\x91\x00`\xe0\x10\x8c\x01\x18<hYc\xd2(\xe3\xa9e\xf4%<C\xa1\x9b\x07\xd0EO\x8a\xd6\xb4\x88\xc2jZ\xed\x14>M\x93\xb2f\xc0\x16{\x02f\xaf\x9fE\x0b\xb5;\xfa\x1e\x83rL\xed\xd9\xbe\x19\xa9S\xeaZ#)\xb4\xdaT\xd2\xda\xf4\xad3\x0b[\xd4)\xee\x8a\xd6\xff\xfb\x14\xc4T\x80C\xe0\x03\t\x00\x84@\x00{\x97\xe0\xd4\x10\x89\xf9\xebh\xda\xaeW_\xd0\xde\x85\x1a\x00\x02\xc0p>\xf7\xb9\x8a\xb5\xd6P\xe9\xdahmJ\xe3\xcb\xa5zhmV*\xdb\x18\x07\x80\x02P\x004\xaa+x\xf7\xa5\xe3]k\xf70\x03\xad\x95\xae\xa8\xa6\xd2\x92o\xdd(n\x08\x89A\xe4\xaa\xd0\x80\x01\x00\xed$\xff\xfb\x14\xc4[\x00C\xb4\x01\t\x00\x84I\x80\x84\x80! \x10\x8dp\x82\xab\xaa\x13,\x86\x84C\x14?J\x90\x86\xa8\xda\xbf\xa7\x150\xa3\xa2\xc8:n.\xb2\x84\x00\x10i{Uz\x9fc/\'j\xc7\xe9U^\xaaU\xd8\x99\xe3\xeeA\xbb\xc7\x14z\xaa\xd0e\xc9\x81\x00\n5kCTX\\\x81,V\xd0%\xda\xed\xea\xa7\xff\xfb\x14\xc4a\x00Ch\x03\r\x00\x04`\x00e\x00"p\x10\x8dp\x7f`w\xb8;s\xd8\xa3\x02\x04\x80\x16\rb\x004\x9dF\xafJ\x1f\x1e\xf5g\\\xadE)\xb4\xc6\xd6\\\xbd\xfa\x90\xc7\xd3\x85\xcd\xe5\x90v\xb9\x82Ha\xd5\xca\xb1;\xa5\xf9\xb6\xfa\xbc\xff\x8a\x1fA\x89i\x86""\xe1N7\xb3(\rJ\xd8\xc4\xbd\xef\xff\xfb\x14\xc4l\x00C\x80\x01\x0b\x00\x84@\x00|\x80a\x10\x10\x8c\x00\x10\x8e?@\xa3=\xbaV\xefG\xddF\xe4\xbc\xc3\xf1\xaaV\xa0\x90\x00<\xe2J<\xc7\xaa\xa8\xd4/\xaer\xb6c\x90\x1a0K\rt-wi\x02%\x89I\x00\xa0\x06\xa7\xb0\x89\x17\xa6|mC\xa3\x9dA\x16l[\x0f\x1f\x87\xb5\xd8\xd4@\xad\xbf\x03\x96\xff\xfb\x14\xc4s\x80\x83\xf0\x03\x08\xc0\x8c`\x00o\x00a`\x10\x8c\x00\x05L\xcd\xd5\x80\x00\x10\x00\x16p\xa2\xd4\xd6\xd8|\x9d\xcf\xaa\xc0\xd9\x05\xab*\xd5_o6\xf2\xf1\xd0=LK\x89\x81\xd2\x0b"\x80\x90\xc6\x0b%\xc7\x1drlgt\xe7f\x92\x97\xb0\xb54\x0b\xfa\xe5*<\x95\xaa\x92F\x92\x8e8\x93$\x00F\x91##\xff\xfb\x14\xc4{\x00C\x90\x01\r \x84I\x80u\x80al\x10\x8c\x00=z\xf5o\x96\xdd\xb6\xffz\xd1\x91\xb3\x8a\xf9k$`\x16\x95\xcc\x80\x12\x1eDbT\x85\xd9Y\x10\x99\x96\xecv.\x9a\x02\xad\xfd\x9f\xd1\x15@\xa9\x95ULAME3.100UUUUUUUUUUUUUUUUUUUU\xff\xfb\x14\xc4\x83\x80\x038\x03\r\x00\x80\x00\x00r\x80"p0\x8d0UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU\xff\xfb\x14\xc4\x8d\x80\xc3\xb4\x01\x0b\x00\x84K\x80~\x00a \x10\x8c\x00UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU\xff\xfb\x14\xc4\x94\x00C\xfc\x03\t\x00\x04@\x00i\x00ad\x00\x8c\x00UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU\xff\xfb\x14\xc4\x9c\x00\x03P\x03\x13\xa0\x84`\x00z\x80\x1e0\x11\x8c\x00UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU'

def log_time(f):
    print(time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime()), file=f, end='')
    print(": ", file=f, end='')

def log_args(f, args):
    print(str(args), file=f)

def log_open(program_args=None):
    log_file = f'{app_identifier}.log'
    log_folder = os.path.join(os.path.expanduser('~'), '.cache', app_identifier)
    os.makedirs(log_folder, exist_ok=True)
    f = open(os.path.join(log_folder, log_file), 'a')
    if program_args:
        log_time(f)
        log_args(f, program_args)
    return f

def log(f, text):
    print(text, end='')
    print(text, file=f, end='')

def log_close(f, print_footer=True):
    if f:
        try:
            if print_footer:
                print()
                print(file=f)
            f.close()
        except:
            pass

def get_session_folder():
    session_folder = os.path.join(os.path.expanduser('~'), '.cache', app_identifier)
    os.makedirs(session_folder, exist_ok=True)
    return session_folder

def get_session_path(session):
    session_file = f'{app_identifier}_{session}.session'
    session_folder = get_session_folder()
    return os.path.join(session_folder, session_file)

def get_user_prompt_from_session(session_file):
    prompt = ""
    with open(session_file, 'r') as f:
        for line in f:
            try:
                message = json.loads(line)
                if message.get('role', '') == 'user':
                    prompt = message.get('content', '')
                    break
            except:
                continue

    return prompt

def list_sessions():
    prefix = f'{app_identifier}_'
    session_folder = get_session_folder()
    # sort by creation date
    sessions = [f for f in os.listdir(session_folder) if f.startswith(prefix) and f.endswith('.session')]
    sessions.sort(key=lambda f: os.path.getctime(os.path.join(session_folder, f)))
    # return name and first line of file
    sessions = [(f[len(prefix):-len('.session')], get_user_prompt_from_session(os.path.join(session_folder, f))) for f in sessions]
    return [(f[0] + '\t' + f[1].strip().replace('\n', ' ')[:80] + '...') for f in sessions]

def session_complete(prefix, parsed_args, **kwargs):
    return [s for s in list_sessions() if s.startswith(prefix)]

def list_models():
    global models
    if not models:
        for provider in ai_providers:
            for model in provider.list_models(get_session_folder()):
                models[model] = provider
    return list(models.keys())

def list_tts_models():
    global tts_models
    if not tts_models:
        for provider in tts_providers:
            for model in provider.list_models(get_session_folder()):
                tts_models[model] = provider
    return list(tts_models.keys())

def list_tts_voices(tts_provider):
    voices = tts_provider.list_voices(get_session_folder())
    return  [ tts_provider.get_voice_name(v) for v in voices ]

def model_complete(prefix, parsed_args, **kwargs):
    models = list_models()
    return [m for m in models if m.startswith(prefix)]

def tts_model_complete(prefix, parsed_args, **kwargs):
    models = list_tts_models()
    return [m for m in models if m.startswith(prefix)]

def switch_to_latest_model(model):
    models = list_models()
    # get the first model which starts with the passed argument, or the leave as is
    latest = next((m for m in models if m.startswith(model)), model)
    if latest != model:
        print_verbose("Model", model, "->", latest)
    return latest

def get_session(session):
    session_path = get_session_path(session)
    messages = []
    if os.path.exists(session_path):
        with open(session_path, 'r') as f:
            for line in f:
                messages.append(json.loads(line))
    return messages

def append_session(session, messages):
    session_path = get_session_path(session)
    with open(session_path, 'a') as f:
        for message in messages:
            print(json.dumps(message), file=f)

def run_stt(audio_file):
    print_verbose("Transcribe", str(audio_file))
    message = default_stt_provider.speech_to_text(model=None, audio_file=audio_file)
    print_verbose("Transcript", str(message))
    return message

def run_tts(tts_provider, tts_model, tts_voice, text, target):
    print_verbose("TTS", str(target), str(text))
    try:
        if isinstance(target, str):
            tts_provider.text_to_speech(text, model=tts_model, voice_id=tts_provider.get_voice_id(tts_voice), speed=1.0, audio_file=target)
        else:
            tts_provider.text_to_speech_stream(text, model=tts_model, voice_id=tts_provider.get_voice_id(tts_voice), speed=1.0, virtual_audio_file=target)
    except Exception as e:
        print_error("Failed to run TTS:", e)
        if isinstance(target, ByteQueueFile): target.close()

def tts(tts_provider, tts_model, tts_voice, text, command, delay_ms, tts_threads = None, tts_queue = None):
    target = None
    if command:
        target = tempfile.mktemp(suffix=".mp3", prefix=f"{app_identifier}.tmp.", dir="/tmp")
    else:
        target = ByteQueueFile()

    if tts_threads is not None and tts_queue is not None:
        print_verbose("Status", "Adding audio data to the TTS queue:", str(target))
        tts_queue.put(target)
        t = threading.Thread(target=run_tts, args=(tts_provider, tts_model, tts_voice, text, target))
        tts_threads.append(t)
        t.start()
    else:
        run_tts(tts_provider, tts_model, tts_voice, text, target)
        handle_audio_data(tts_provider, target, command, delay_ms)

def optimize_text_for_tts(text, optimize):
    if optimize:
        before = text
        text, lang = optimize_for_tts(text)
        print_verbose("Optimize", f'"{before}" -> "{text}"')

    return text

def extract_sentence(text, min_len=MIN_SENTENCE_LENGTH, max_len=MAX_SENTENCE_LENGTH) -> tuple[str, str] | None:
    text_len = len(text)

    match = re.search(rf'^.{{{min_len},{max_len}}}[^0-9](?:[.!?]{{1,3}})(?:\s|$)', text, re.DOTALL)
    if not match:
        match = re.search(rf'^.{{{min_len},{max_len}}}[^0-9](?:[-—,;:])(?:\s|$)', text, re.DOTALL)

    if text_len <= max_len:
        # just return the whole text if it's shorter than max_len
        segment_end = text_len
    else:
        # otherwise try with new line or resort to a space
        if not match:
            match = re.search(rf'^.{{{min_len},{max_len}}}\n', text, re.DOTALL)
        if not match:
            match = re.search(rf'^.{{{min_len},{max_len}}}\s', text, re.DOTALL)

        if match:
            segment_end = match.end()
        else:
            # as a last resort, just take the last max_len characters
            print(f'WARNING: Unable to extract sentence, will split text at {max_len} characters', file=sys.stderr, flush=True)
            segment_end = max_len

    segment = text[:segment_end].strip()
    remainder = text[segment_end:].strip()
    return segment, remainder

def play_silence_to_keep_audio_alive():
    global LAST_PLAYED_TIMESTAMP
    now = int(time.time())
    if now > LAST_PLAYED_TIMESTAMP + AUDIO_TIMEOUT_SECONDS:
        print_verbose("Play", f'Playing 500ms silence')
        def memory_stream(soundfile: miniaudio.DecodedSoundFile) -> miniaudio.PlaybackCallbackGeneratorType:
            required_frames = yield b"" # generator initialization
            current = 0
            samples = memoryview(soundfile.samples) # avoid needless memory copying
            while current < len(samples):
                sample_count = required_frames * soundfile.nchannels
                output = samples[current:current + sample_count]
                current += sample_count
                required_frames = yield output

        decoded = miniaudio.mp3_read_s16(SILENCE_500_BYTES)
        stream = memory_stream(decoded)
        next(stream)  # start the generator
        with miniaudio.PlaybackDevice(nchannels=1) as device:
            device.start(stream)
            while device.callback_generator is not None:
                time.sleep(0.1)

        if LAST_PLAYED_TIMESTAMP == 0:
            LAST_PLAYED_TIMESTAMP = int(time.time())

        print_verbose("Play", "Done playing silence")

def play_audio_file(tts_provider, audio_file, delay_ms):
    global LAST_PLAYED_TIMESTAMP
    try:
        miniaudio_format = miniaudio.FileFormat.UNKNOWN
        miniaudio_sample_format = miniaudio.SampleFormat.UNKNOWN
        tts_format = tts_provider.format()
        if tts_format:
            if tts_format == TTSProvider.Format.MP3:
                miniaudio_format = miniaudio.FileFormat.MP3
            elif tts_format == TTSProvider.Format.WAV:
                miniaudio_format = miniaudio.FileFormat.WAV

        tts_dtype = tts_provider.dtype()
        if tts_dtype:
            tts_dtype = tts_dtype.lower()
            if tts_dtype == 'float32':
                miniaudio_sample_format = miniaudio.SampleFormat.FLOAT32
            elif tts_dtype == 'signed16':
                miniaudio_sample_format = miniaudio.SampleFormat.SIGNED16

        if isinstance(audio_file, ByteQueueFile):
            if audio_file.closed:
                delay_ms = 0
            else:
                play_silence_to_keep_audio_alive()
                if audio_file.closed:
                    delay_ms = 0

        print_verbose("Play", f"{audio_file} with delay {delay_ms}ms, format: {miniaudio_format}, output: {miniaudio_sample_format}")
        stream = miniaudio.stream_any(audio_file, source_format=miniaudio_format, output_format=miniaudio_sample_format, nchannels=tts_provider.channels(), sample_rate=tts_provider.samplerate(), frames_to_read=tts_provider.blocksize())
        try:
            with miniaudio.PlaybackDevice(output_format=miniaudio_sample_format, nchannels=tts_provider.channels(), sample_rate=tts_provider.samplerate(), buffersize_msec=int(delay_ms)) as device:
                device.start(stream)
                while device.callback_generator is not None:
                    time.sleep(0.2)
        finally:
            LAST_PLAYED_TIMESTAMP = int(time.time())
            stream.close()
    except KeyboardInterrupt:
        sys.exit(0)
    except Exception as e:
        print('Failed to play audio:', str(e), file=sys.stderr, flush=True)

def handle_audio_file(tts_provider, target, command, delay_ms):
    if command:
        if not isinstance(target, str):
            print_error("Unexpected argument for target:", target)
            return

        if '{}' in command:
            command = command.format(target)
        if delay_ms > 0:
            print_verbose("Status", f'Delaying audio by {delay_ms}ms')
            time.sleep(delay_ms / 1000.0)
        print_verbose("Status", "Running audio command:", str(command))
        try:
            subprocess.run(['bash', '-c', command], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except KeyboardInterrupt:
            sys.exit(0)
        except Exception as e:
            print_error("Failed to run output audio command", str(target), ':', e)
    else:
        try:
            print_verbose("Status", "Playing audio data: ", str(target))
            play_audio_file(tts_provider, target, delay_ms)
        except KeyboardInterrupt:
            sys.exit(0)
        except Exception as e:
            print_error("Failed to play audio", str(target), ':', e)

def run_audio_file_queue(tts_provider, command, delay_ms, command_queue, stop_event):
    try:
        print_verbose("Status", "Starting audio queue")
        running = True
        while running:
            if stop_event.is_set():
                print_verbose("Status", "Stopping audio queue")
                running = False

            while not command_queue.empty():
                target = command_queue.get()
                print_verbose("Status", "Handling audio data", str(target))
                handle_audio_file(tts_provider, target, command, delay_ms)

            time.sleep(0.1)
    except KeyboardInterrupt:
        sys.exit(0)
    except Exception as e:
        print_error("Failed to run audio queue:", e)

def handle_audio_data(tts_provider, target, command, delay_ms, command_queue = None):
    if command_queue:
        print_verbose("Status", "Adding audio data from the TTS queue to the command queue:", str(target))
        command_queue.put(target)
    else:
        handle_audio_file(tts_provider, target, command, delay_ms)

def merge_content(file, audio_file, text, prompt = False):
    result = None
    if text:
        if isinstance(text, str):
            result = text
        else:
            result = ' '.join(text)

    if prompt:
        print('(Press Ctrl-D end input.)', file=sys.stderr)
        lines = []
        while True:
            try:
                line = input()
            except KeyboardInterrupt:
                print('')
                exit(1)
            except EOFError:
                break

            lines.append(line)

        if len(lines) > 0:
            if result and len(result) > 0:
                result += '\n'
            else:
                result = ''

            result += '\n'.join(lines)

    audio_content = None
    if audio_file:
        audio_content = run_stt(audio_file)

    if audio_content:
        if result and len(result) > 0:
            result += '\n'
        else:
            result = ''

        result += audio_content

    file_content = None
    if file == '-':
        file_content = sys.stdin.read()
    elif file:
        with open(file) as f: file_content = f.read()

    if file_content:
        if result and len(result) > 0:
            result += '\n'
        else:
            result = ''

        result += file_content

    return result

def main():
    parser = argparse.ArgumentParser(description='Chat based AI assistant', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("--version", action="version", version=f"{os.path.basename(sys.argv[0])} {VERSION}")
    parser.add_argument('-m', '--model', type=str, help='Which model to use', default='perplexity').completer = model_complete
    parser.add_argument('-t', '--tts-model', type=str, help='Which TTS model to use', default='tts-1').completer = tts_model_complete
    parser.add_argument('-v', '--tts-voice', type=str, help='Which TTS voice to use', required=False)
    parser.add_argument('-f', '--file', type=str, help='File with content to append to the prompt, - for stdin', default='').completer = lambda: [f for f in os.listdir('.') if os.path.isfile(f)]
    parser.add_argument('-i', '--input', action='store_true', help='Prompt for multi line input', default=False)
    parser.add_argument("-p", "--print-prompt", action="store_true", help="Print the prompt", default=False)
    parser.add_argument('-a', '--audio', type=str, help='Audio file with content to append to the prompt (will use speech-to-text to transcribe)', default='').completer = lambda: [f for f in os.listdir('.') if os.path.isfile(f)]
    parser.add_argument('-o', '--output', type=str, help='Output format', choices=['text', 'audio', 'audio+text'], default='text')
    parser.add_argument('-c', '--output-audio-command', type=str, help='Output command for audio files, e.g "mpg123 -q {}"', default='')
    parser.add_argument('-d', '--output-audio-delay-ms', type=int, help='Output audio delay in milliseconds when streaming TTS audio', default=200)
    parser.add_argument('-r', '--role', type=str, help='Which role to take')
    parser.add_argument('--role-file', type=str, help='File with content to append to the role, - for stdin', default='').completer = lambda: [f for f in os.listdir('.') if os.path.isfile(f)]
    parser.add_argument('-w', '--wait', action='store_true', help='Wait for the full response, don\'t stream', default=False)
    parser.add_argument('-s', '--session', type=str, help='Session to start or reuse', default='').completer = session_complete
    parser.add_argument("--verbose", action="store_true", help="Print details like used model and session", default=False)
    parser.add_argument("--no-session", action="store_true", help="Prevent session creation and always start fresh", default=False)
    parser.add_argument("--no-model-switch", action="store_true", help="Prevent switching to latest model", default=False)
    parser.add_argument("--no-sources", action="store_true", help="Don't print sources, e.g. internet sources", default=False)
    parser.add_argument("--no-tts-optimization", action="store_true", help="Don't optimize TTS, e.g. replace numbers with words", default=False)
    list_group = parser.add_mutually_exclusive_group()
    list_group.add_argument('--list-sessions', action='store_true', help='List sessions', default=False)
    list_group.add_argument('--list-models', action='store_true', help='List models', default=False)
    list_group.add_argument('--list-tts-models', action='store_true', help='List TTS models', default=False)
    list_group.add_argument('--list-tts-voices', action='store_true', help='List TTS voices of the selected TTS model', default=False)
    list_group.add_argument('--print-session', action='store_true', help='Print session', default=False)
    parser.add_argument('prompt', type=str, nargs='*', help='The prompt to send')
    argcomplete.autocomplete(parser)

    args = parser.parse_args()

    if args.print_session:
        if not args.session:
            print_error("No session specified")
            parser.print_help()
            sys.exit(1)
        messages = get_session(args.session)
        if not messages:
            print_error("Session not found")
            sys.exit(1)

        first = True
        for message in messages:
            if first:
                first = False
            else:
                print()
            role = message['role']
            role = role[0].upper() + role[1:]
            print(role + ':\n' + message['content'])
        sys.exit(0)

    if args.list_sessions:
        for session in list_sessions():
            print(session)
        sys.exit(0)

    use_sessions = not args.no_session

    f = None
    try:
        utils.VERBOSE = args.verbose

        global ai_providers
        ai_providers = get_ai_providers()

        global tts_providers
        tts_providers = get_tts_providers()

        global stt_providers, default_stt_provider
        stt_providers = get_stt_providers()
        default_stt_provider = stt_providers[0]

        if args.list_models:
            for model in list_models():
                print(model)
            sys.exit(0)

        if args.list_tts_models:
            for model in list_tts_models():
                print(model)
            sys.exit(0)

        if args.tts_model and not args.tts_model in list_tts_models():
            print_error("Unknown TTS model:", "'" + args.tts_model + "'.\nAvailable models:\n" + '\n'.join(list_tts_models()))
            sys.exit(1)

        if args.list_tts_voices:
            if not args.tts_model:
                print_error("No TTS model specified")
                parser.print_help()
                sys.exit(1)
            for voice in list_tts_voices(tts_models[args.tts_model]):
                print(voice)
            sys.exit(0)

        if not args.prompt and not args.file and not args.audio and not args.input:
            print_error("No prompt or file specified")
            parser.print_help()
            sys.exit(0)

        if not args.model in list_models():
            print_error("Unknown model:", "'" + args.tts_model + "'.\nAvailable models:\n" + '\n'.join(list_models()))
            sys.exit(1)

        ai_provider = models[args.model]
        tts_provider = tts_models[args.tts_model]
        tts_voice = tts_provider.get_voice_by_name(args.tts_voice, get_session_folder())
        print_verbose("Provider", ai_provider.name())
        print_verbose("Model", args.model)
        print_verbose("TTS provider", tts_provider.name())
        print_verbose("TTS model", args.tts_model)
        voice_name = tts_provider.get_voice_name(tts_voice)
        voice_id = tts_provider.get_voice_id(tts_voice)
        if voice_name != voice_id:
            print_verbose("TTS voice", f'{voice_name} ({voice_id})')
        else:
            print_verbose("TTS voice", voice_name)

        # make sure we use the latest model
        if not args.no_model_switch:
            args.model = switch_to_latest_model(args.model)

        if use_sessions:
            # make sure the AI provider supports sessions
            use_sessions = ai_provider.supports_sessions()

        f = log_open(args)

        messages = []

        session = None
        if use_sessions:
            if args.session:
                session = args.session
                print_verbose("Session", session)
            elif sys.__stdin__.isatty():
                # use pid of parent bash shell as session when in interactive shell
                proc = psutil.Process(os.getpid())
                while proc is not None:
                    if len(proc.cmdline()) == 1 and proc.cmdline()[0].endswith('bash') and proc.exe().endswith('/bin/bash'):
                        # use parent bash shell pid as session
                        timestamp = int(proc.create_time())
                        session = str(proc.pid) + "_" + time.strftime("%Y-%m-%d_%H-%M-%S", time.gmtime(timestamp))

                        print_verbose("Session", str(session), "(tty)")
                        break
                    proc = proc.parent()

        if session:
            # load messages from session
            session_messages = get_session(session)
            messages.extend(session_messages)
        else:
            print_verbose("Session", "N/A")

        new_messages = []
        role = merge_content(args.role_file, None, args.role)
        if role:
            # don't change role
            for message in messages:
                if message["role"] == "system":
                    # check if args.role is different
                    if role != message["content"]:
                        print_error("Role cannot be changed in an existing session, before:", message["content"], ", after:", role)
                        parser.print_help()
                        sys.exit(0)
                    break

            new_messages.append({"role": "system", "content": role})

        raw_prompt = merge_content(args.file, args.audio, args.prompt, args.input)
        if raw_prompt is None:
            print_error("Prompt is empty")
            parser.print_help()
            sys.exit(0)

        if args.print_prompt:
            print(raw_prompt)

        new_messages.append({"role": "user", "content": raw_prompt})
        messages.extend(new_messages)

        print_verbose("Messages", str(len(messages)))

        sources = None
        if args.no_sources == False:
            sources = OrderedDict()

        result = ai_provider.chat_completion(messages, args.model, not args.wait)

        handle_metadata_func = None
        if utils.VERBOSE:
            handle_metadata_func = print_verbose
            print_verbose()

        stop_event = None
        command_queue = None
        command_thread = None
        if not args.wait and (args.output == 'audio' or args.output == 'audio+text'):
            stop_event = threading.Event()
            command_queue = queue.Queue()
            command_thread = threading.Thread(target=run_audio_file_queue, args=(tts_provider, args.output_audio_command, int(args.output_audio_delay_ms), command_queue, stop_event))
            command_thread.start()

        tts_threads = []
        tts_queue = queue.Queue()

        if args.output == 'audio' or args.output == 'audio+text':
            # to make sure that bluetooth audio is on, we play silence first
            play_silence_to_keep_audio_alive()

        answer = ""
        if args.wait:
            answer = ai_provider.convert_result_to_text(result, sources, handle_metadata_func)
            if not sources:
                answer = ai_provider.remove_source_references(answer)

            if utils.VERBOSE:
                print_verbose()

            if args.output == 'audio' or args.output == 'audio+text':
                answer = optimize_text_for_tts(answer, not args.no_tts_optimization)
                if len(answer) > MAX_SENTENCE_LENGTH:
                    while len(answer) > 0:
                        sentence, answer = extract_sentence(answer, max_len=int(MAX_SENTENCE_LENGTH / 2))
                        tts(tts_provider, args.tts_model, tts_voice, sentence, args.output_audio_command, int(args.output_audio_delay_ms))
                elif len(answer) > 0:
                    tts(tts_provider, args.tts_model, tts_voice, answer, args.output_audio_command, int(args.output_audio_delay_ms))

            if args.output == 'text' or args.output == 'audio+text':
                log(f, answer)
        else:
            first = True
            first_tts = True
            answer = ""
            segment = ""
            text_chunks = []
            for chunk in result:
                answer_chunk = ai_provider.convert_chunk_to_text(chunk, text_chunks, sources, handle_metadata_func)
                if not sources:
                    answer_chunk = ai_provider.remove_source_references(answer_chunk)

                if first:
                    first = False
                    handle_metadata_func = None
                    print_verbose()

                if answer_chunk is not None:
                    answer += answer_chunk
                    segment += answer_chunk
                    if args.output == 'text' or args.output == 'audio+text':
                        log(f, answer_chunk)
                    if args.output == 'audio' or args.output == 'audio+text':
                        # make sure we try to send sentences to tts
                        max_len = MAX_SENTENCE_LENGTH
                        if first_tts:
                            first_tts = False
                            # optimization to start audio more quickly
                            max_len = int(MAX_SENTENCE_LENGTH / 2)
                        sentence, segment = extract_sentence(segment, max_len=max_len)
                        if sources:
                            # always remove source refs for TTS, if they are requested
                            sentence = ai_provider.remove_source_references(sentence).strip()
                        if sentence:
                            if not tts_queue.empty():
                                handle_audio_data(tts_provider, tts_queue.get(), args.output_audio_command, int(args.output_audio_delay_ms), command_queue)

                            sentence = optimize_text_for_tts(sentence, not args.no_tts_optimization)
                            # TODO: might be too large, check tts_provider and cut again
                            tts(tts_provider, args.tts_model, tts_voice, sentence, args.output_audio_command, int(args.output_audio_delay_ms), tts_threads, tts_queue)

        if sources:
            print()
            for k, v in sources.items():
                title = str(v).strip()
                url = str(k).strip()
                if len(url) == 0 and len(title) == 0:
                    continue

                if len(title) > 0:
                    log(f, f'\n{title}\n{url}\n')
                else:
                    log(f, f'{url}\n')

        if session:
            # save query and answer to session
            new_messages.append({"role": "assistant", "content": answer})
            append_session(session, new_messages)

        # wait for tts to finish
        if not args.wait and args.output == 'audio' or args.output == 'audio+text':
            # check if there is a segment left
            segment = segment.strip()
            if sources:
                # always remove source refs for TTS, if they are requested
                segment = ai_provider.remove_source_references(segment).strip()

            if segment:
                segment = optimize_text_for_tts(segment, not args.no_tts_optimization)
                while len(segment) > 0:
                    sentence, segment = extract_sentence(segment)
                    tts(tts_provider, args.tts_model, tts_voice, sentence, args.output_audio_command, int(args.output_audio_delay_ms), tts_threads, tts_queue)

            # wait for all tts tasks to be done
            print_verbose("Status", "Waiting for TTS threads to finish...")
            sys.stdout.flush()
            try:
                for t in tts_threads:
                    t.join()
            except KeyboardInterrupt:
                sys.exit(0)

            print_verbose("Status", "Handling remaining audio files...")
            while not tts_queue.empty():
                handle_audio_data(tts_provider, tts_queue.get(), args.output_audio_command, int(args.output_audio_delay_ms), command_queue)

            if command_thread and stop_event:
                print_verbose("Status", "Waiting for command thread to finish...")
                sys.stdout.flush()
                stop_event.set()
                try:
                    command_thread.join()
                except KeyboardInterrupt:
                    sys.exit(0)

    except KeyboardInterrupt:
        sys.exit(0)
    except Exception as e:
        print_error("Failed to run inference:", e)
        log_close(f, False)
        sys.exit(1)
    finally:
        if ai_providers:
            for p in ai_providers:
                try:
                    p.close()
                except Exception as e:
                    print_error("Failed to close AI provider", p.name(), ":", e)
        if tts_providers:
            for p in tts_providers:
                try:
                    p.close()
                except Exception as e:
                    print_error("Failed to close TTS provider", p.name(), ":", e)


    log_close(f)


if __name__ == "__main__":
    main()