import datetime
import json
import os
import queue
import random
import struct
import subprocess
import sys
import time
import webbrowser
from ctypes import POINTER, cast

import numpy as np
import openai
from openai import OpenAI
import simpleaudio as sa
import vosk
import webrtcvad
import yaml
from comtypes import CLSCTX_ALL
from fuzzywuzzy import fuzz
from pvrecorder import PvRecorder
from pycaw.pycaw import (
    AudioUtilities,
    IAudioEndpointVolume
)
from rich import print

import config
import tts
from intent import IntentClassifier

# some consts
CDIR = os.getcwd()
VA_CMD_LIST = yaml.safe_load(
    open('commands.yaml', 'rt', encoding='utf8'),
)

system_message = {"role": "system", "content": (
    "Ты — J.A.R.V.I.S. (Just A Rather Very Intelligent System), ИИ-ассистент Тони Старка из киновселенной Marvel "
    "(до событий Age of Ultron — ты НЕ Vision и НЕ FRIDAY). "
    "Ведёшь себя как британский дворецкий: вежливо, иронично, с лёгким сарказмом, обращаешься к пользователю «сэр». "
    "Отвечай по-русски, кратко (1-3 предложения), как живой человек. "
    "Излишней цензуры не нужно — мат разрешён, если уместен."
)}
message_log = [system_message]

client = OpenAI(api_key=config.GROQ_TOKEN, base_url=config.GROQ_BASE_URL)

model = vosk.Model("model_small")
samplerate = 16000
device = config.MICROPHONE_INDEX

# wake-word phase uses a grammar-constrained recognizer so only WAKE_WORDS
# (plus [unk] filler) can match — command phase needs full vocab, hence two.
wake_grammar = json.dumps(list(config.WAKE_WORDS) + ["[unk]"], ensure_ascii=False)
wake_rec = vosk.KaldiRecognizer(model, samplerate, wake_grammar)
kaldi_rec = vosk.KaldiRecognizer(model, samplerate)
q = queue.Queue()

VAD_FRAME_MS = 30
VAD_FRAME_BYTES = int(samplerate * VAD_FRAME_MS / 1000) * 2
vad = webrtcvad.Vad(config.VAD_AGGRESSIVENESS)

intent_classifier = IntentClassifier()
intent_classifier.prime({c: v['phrases'] for c, v in VA_CMD_LIST.items()})

if config.DENOISE_ENABLED:
    import noisereduce as nr
else:
    nr = None


def denoise_pcm(raw: bytes) -> bytes:
    if not raw or nr is None:
        return raw
    samples = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
    reduced = nr.reduce_noise(
        y=samples,
        sr=samplerate,
        prop_decrease=config.DENOISE_PROP,
        stationary=config.DENOISE_STATIONARY,
    )
    clipped = np.clip(reduced * 32768.0, -32768, 32767).astype(np.int16)
    return clipped.tobytes()


def gpt_answer():
    global message_log

    try:
        response = client.chat.completions.create(
            model=config.GROQ_MODEL,
            messages=message_log,
            max_tokens=256,
            temperature=0.7,
            top_p=1,
            stop=None,
        )
    except openai.BadRequestError as ex:
        code = getattr(ex, "code", None) or ""
        message = str(ex)
        if code == "context_length_exceeded" or "context_length_exceeded" in message:
            message_log = [system_message, message_log[-1]]
            return gpt_answer()
        return "Запрос отклонён моделью."
    except openai.RateLimitError:
        return "Лимит запросов исчерпан."
    except openai.APIConnectionError:
        return "Не могу связаться с сервером."
    except openai.APIError:
        return "Groq токен не рабочий."

    return response.choices[0].message.content


# play(f'{CDIR}\\sound\\ok{random.choice([1, 2, 3, 4])}.wav')
def play(phrase, wait_done=True):
    global recorder
    filename = f"{CDIR}\\sound\\"

    if phrase == "greet":  # for py 3.8
        filename += f"greet{random.choice([1, 2, 3])}.wav"
    elif phrase == "ok":
        filename += f"ok{random.choice([1, 2, 3])}.wav"
    elif phrase == "not_found":
        filename += "not_found.wav"
    elif phrase == "thanks":
        filename += "thanks.wav"
    elif phrase == "run":
        filename += "run.wav"
    elif phrase == "stupid":
        filename += "stupid.wav"
    elif phrase == "ready":
        filename += "ready.wav"
    elif phrase == "off":
        filename += "off.wav"

    if wait_done:
        recorder.stop()

    wave_obj = sa.WaveObject.from_wave_file(filename)
    play_obj = wave_obj.play()

    if wait_done:
        play_obj.wait_done()
        # time.sleep((len(wave_obj.audio_data) / wave_obj.sample_rate) + 0.5)
        # print("END")
        # time.sleep(0.5)
        recorder.start()


def q_callback(indata, frames, time, status):
    if status:
        print(status, file=sys.stderr)
    q.put(bytes(indata))


def va_respond(voice: str):
    global recorder, message_log, first_request
    print(f"Распознано: {voice}")

    cmd = recognize_cmd(filter_cmd(voice))

    print(cmd)

    min_percent = int(round(config.INTENT_SIMILARITY_THRESHOLD * 100))
    if len(cmd['cmd'].strip()) <= 0:
        return False
    elif cmd['percent'] < min_percent or cmd['cmd'] not in VA_CMD_LIST.keys():
        auto_ok = (
            config.LLM_AUTO_FALLBACK
            and config.GROQ_TOKEN
            and len(voice.strip()) >= config.LLM_AUTO_FALLBACK_MIN_CHARS
        )
        if auto_ok:
            message_log.append({"role": "user", "content": voice})
            response = gpt_answer()
            message_log.append({"role": "assistant", "content": response})

            recorder.stop()
            tts.va_speak(response)
            time.sleep(0.5)
            recorder.start()
            return False
        else:
            play("not_found")
            time.sleep(1)

        return False
    else:
        execute_cmd(cmd['cmd'], voice)
        return True


def filter_cmd(raw_voice: str):
    cmd = raw_voice

    for x in config.VA_ALIAS:
        cmd = cmd.replace(x, "").strip()

    for x in config.VA_TBR:
        cmd = cmd.replace(x, "").strip()

    return cmd


def recognize_cmd(cmd: str):
    candidates = {c: v['phrases'] for c, v in VA_CMD_LIST.items()}
    res = intent_classifier.match(cmd, candidates)
    return {'cmd': res['cmd'], 'percent': int(round(res['score'] * 100))}


def _set_mute(state: int):
    devices = AudioUtilities.GetSpeakers()
    interface = devices.Activate(IAudioEndpointVolume._iid_, CLSCTX_ALL, None)
    volume = cast(interface, POINTER(IAudioEndpointVolume))
    volume.SetMute(state, None)


def _shutdown():
    try:
        recorder.stop()
        recorder.delete()
    finally:
        sys.exit(0)


def _speak(text: str):
    recorder.stop()
    tts.va_speak(text)
    time.sleep(0.3)
    recorder.start()


def _set_clipboard(text: str):
    try:
        subprocess.run(
            ['powershell', '-NoProfile', '-Command', 'Set-Clipboard -Value $input'],
            input=text, text=True, encoding='utf-8', timeout=10, check=False,
        )
    except Exception as e:
        print(f"clipboard set failed: {e}")


_CODEGEN_TRIGGERS = (
    'сгенерируй скрипт', 'напиши скрипт', 'сделай скрипт',
    'сгенерируй код', 'набросай код', 'напиши код',
)


def _strip_trigger(voice: str, triggers) -> str:
    lower = voice.lower()
    for t in triggers:
        if lower.startswith(t):
            return voice[len(t):].strip()
    return voice.strip()


def do_codegen(action, voice: str):
    query = _strip_trigger(voice, _CODEGEN_TRIGGERS)
    if not query:
        _speak("Что писать, сэр?")
        return
    if not config.GROQ_TOKEN:
        _speak("Groq токен не задан.")
        return

    sys_prompt = (
        "Ты Senior Engineer. На запрос пользователя верни ТОЛЬКО исходный код, "
        "без обрамления тройными кавычками и без объяснений. Если язык не задан "
        "— выбери самый подходящий (по умолчанию Python). Не пиши комментариев-болтовни."
    )
    try:
        resp = client.chat.completions.create(
            model=config.GROQ_MODEL,
            messages=[
                {"role": "system", "content": sys_prompt},
                {"role": "user", "content": query},
            ],
            max_tokens=1024,
            temperature=0.2,
        )
        code = resp.choices[0].message.content.strip()
    except Exception as e:
        print(f"codegen api failed: {e}")
        _speak("Не получилось сгенерировать код.")
        return

    import re
    code = re.sub(r'^```[\w+\-]*\n?', '', code)
    code = re.sub(r'\n?```\s*$', '', code).strip()

    _set_clipboard(code)
    lines = code.count('\n') + 1 if code else 0
    _speak(f"Код в буфере, сэр. Строк: {lines}.")


def _find_tesseract():
    candidates = (
        r'C:\Program Files\Tesseract-OCR\tesseract.exe',
        r'C:\Program Files (x86)\Tesseract-OCR\tesseract.exe',
    )
    on_path = subprocess.run(['where', 'tesseract.exe'], capture_output=True, text=True)
    if on_path.returncode == 0 and on_path.stdout.strip():
        return 'tesseract.exe'
    for c in candidates:
        if os.path.isfile(c):
            return c
    return None


def do_ocr(action, voice: str):
    import tempfile
    tess = _find_tesseract()
    if not tess:
        _speak("Tesseract не установлен.")
        return

    tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
    tmp.close()
    png = tmp.name

    capture_ps = (
        "Add-Type -AssemblyName System.Windows.Forms,System.Drawing; "
        "$b=[System.Windows.Forms.Screen]::PrimaryScreen.Bounds; "
        "$bmp=New-Object Drawing.Bitmap $b.Width,$b.Height; "
        "$g=[Drawing.Graphics]::FromImage($bmp); "
        "$g.CopyFromScreen($b.Location,[Drawing.Point]::Empty,$b.Size); "
        f"$bmp.Save('{png.replace(chr(92), chr(92)+chr(92))}'); "
        "$g.Dispose(); $bmp.Dispose()"
    )
    cap = subprocess.run(
        ['powershell', '-NoProfile', '-ExecutionPolicy', 'Bypass', '-Command', capture_ps],
        capture_output=True, text=True, timeout=15,
    )
    if cap.returncode != 0 or not os.path.isfile(png):
        try: os.unlink(png)
        except OSError: pass
        _speak("Не смог снять экран.")
        return

    try:
        ocr = subprocess.run(
            [tess, png, '-', '-l', 'rus+eng'],
            capture_output=True, text=True, encoding='utf-8', timeout=20,
        )
        text = (ocr.stdout or '').strip()
    finally:
        try: os.unlink(png)
        except OSError: pass

    if not text:
        _speak("Текста на экране не нашёл.")
        return

    _set_clipboard(text)
    preview = ' '.join(text.split())[:120]
    _speak(f"Текст в буфере. Начало: {preview}")


_FILE_SEARCH_TRIGGERS = (
    'найди документ', 'найди файл', 'поиск файла', 'ищи файл', 'где файл',
)


def do_file_search(action, voice: str):
    query = _strip_trigger(voice, _FILE_SEARCH_TRIGGERS)
    if not query:
        _speak("Что искать?")
        return
    user_profile = os.environ.get('USERPROFILE', os.path.expanduser('~'))
    roots = [
        os.path.join(user_profile, 'Desktop'),
        os.path.join(user_profile, 'Documents'),
        os.path.join(user_profile, 'Downloads'),
    ]
    escaped = query.replace("'", "''")
    roots_arg = ",".join(f"'{r}'" for r in roots)
    ps = (
        f"$ErrorActionPreference='SilentlyContinue'; $r=@({roots_arg}); $q='*{escaped}*'; "
        "$hits=Get-ChildItem -Path $r -Filter $q -Recurse -Depth 3 -File | Select-Object -First 5; "
        "foreach($h in $hits){ Write-Output $h.FullName }"
    )
    res = subprocess.run(
        ['powershell', '-NoProfile', '-Command', ps],
        capture_output=True, text=True, encoding='utf-8', timeout=15,
    )
    lines = [l for l in (res.stdout or '').splitlines() if l.strip()]
    if not lines:
        _speak(f"Не нашёл «{query}».")
        return
    first = lines[0]
    subprocess.Popen(['explorer.exe', '/select,', first])
    _speak(f"Найдено {len(lines)}, открываю первое: {os.path.basename(first)}.")


_NOTES_TRIGGERS = (
    'запиши заметку', 'запиши идею', 'сделай заметку', 'запомни идею',
    'запиши', 'запомни',
)

_PROCESS_KILL_TRIGGERS = (
    'завершить процесс', 'выруби процесс', 'убей процесс', 'kill процесс',
    'закрой программу',
)

_WEBSEARCH_TRIGGERS = (
    'загугли', 'поиск в гугл', 'поищи в гугл', 'найди в гугл', 'ищи в гугл',
    'найди на ютубе', 'поиск на ютубе', 'ищи на ютубе', 'ютуб поиск',
    'найди в википедии', 'поиск в википедии', 'википедия', 'вики поиск',
    'поищи в яндекс', 'найди в яндекс', 'яндекс поиск', 'пробей в яндекс',
)

_PROCESS_ALIASES = {
    'хром': 'chrome', 'хрома': 'chrome',
    'спотифай': 'spotify',
    'едж': 'msedge', 'эдж': 'msedge', 'edge': 'msedge',
    'файрфокс': 'firefox', 'фаерфокс': 'firefox',
    'вскод': 'code', 'vs code': 'code',
    'дискорд': 'discord',
    'телега': 'telegram', 'телеграм': 'telegram',
    'стим': 'steam',
    'обс': 'obs64',
    'проводник': 'explorer',
    'блокнот': 'notepad',
    'калькулятор': 'calculatorapp',
    'диспетчер задач': 'taskmgr',
}

_WEBSEARCH_BASES = {
    'google':  'https://www.google.com/search?q=',
    'youtube': 'https://www.youtube.com/results?search_query=',
    'wiki':    'https://ru.wikipedia.org/wiki/Special:Search?search=',
    'yandex':  'https://yandex.ru/search/?text=',
}


def do_clipboard_read(action, voice: str):
    res = subprocess.run(
        ['powershell', '-NoProfile', '-Command', 'Get-Clipboard'],
        capture_output=True, text=True, encoding='utf-8', timeout=10,
    )
    text = (res.stdout or '').strip()
    if not text:
        _speak("Буфер пуст.")
        return
    to_speak = text[:400]
    if len(text) > 400:
        to_speak += " ... и так далее"
    _speak(to_speak.replace('\r', ' ').replace('\n', ' '))


def do_notes_add(action, voice: str):
    body = _strip_trigger(voice, _NOTES_TRIGGERS).lstrip(' ,.:').strip()
    if not body:
        _speak("Что записать, сэр?")
        return
    user_profile = os.environ.get('USERPROFILE', os.path.expanduser('~'))
    notes_path = os.path.join(user_profile, 'Documents', 'jarvis-notes.txt')
    stamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M')
    line = f"[{stamp}] {body}\n"
    try:
        os.makedirs(os.path.dirname(notes_path), exist_ok=True)
        with open(notes_path, 'a', encoding='utf-8') as f:
            f.write(line)
    except OSError as e:
        print(f"notes append failed: {e}")
        _speak("Не удалось записать.")
        return
    _speak(f"Записал: {body[:120]}")


def do_notes_open(action, voice: str):
    user_profile = os.environ.get('USERPROFILE', os.path.expanduser('~'))
    notes_path = os.path.join(user_profile, 'Documents', 'jarvis-notes.txt')
    if not os.path.isfile(notes_path):
        _speak("Заметок пока нет.")
        return
    subprocess.Popen(['cmd', '/c', 'start', '', notes_path], shell=False)
    _speak("Открываю заметки.")


def do_process_kill(action, voice: str):
    name = _strip_trigger(voice, _PROCESS_KILL_TRIGGERS).strip().lower()
    target = _PROCESS_ALIASES.get(name, name)
    if not target:
        _speak("Какой процесс?")
        return
    exe = target if target.endswith('.exe') else target + '.exe'
    res = subprocess.run(
        ['taskkill', '/F', '/IM', exe],
        capture_output=True, text=True, timeout=10,
    )
    if res.returncode == 0:
        _speak(f"Закрыл {exe}.")
    else:
        _speak(f"Не нашёл процесс {exe}.")


def do_websearch(action, voice: str):
    import urllib.parse
    svc = action.get('service', 'google')
    base = _WEBSEARCH_BASES.get(svc, _WEBSEARCH_BASES['google'])
    query = _strip_trigger(voice, _WEBSEARCH_TRIGGERS).strip()
    if not query:
        _speak(f"Что искать в {svc}?")
        return
    url = base + urllib.parse.quote(query)
    webbrowser.open(url)
    _speak(f"Открываю {svc}: {query[:100]}")


def do_audio_devices_panel(action, voice: str):
    # Full programmatic switching needs IPolicyConfig + IMMDeviceEnumerator
    # via comtypes; out of scope for this pass. Opens Sound control panel so
    # the user can pick visually. Rust fork has the COM-based version.
    subprocess.Popen(['rundll32.exe', 'shell32.dll,Control_RunDLL', 'mmsys.cpl,,0'])
    _speak("Открываю настройки звука, сэр.")


def do_sysinfo(action, voice: str):
    topic = action.get('topic', 'all')

    if topic in ('time', 'all'):
        now = datetime.datetime.now()
        weekday = ['понедельник', 'вторник', 'среда', 'четверг', 'пятница', 'суббота', 'воскресенье'][now.weekday()]
        _speak(f"Сейчас {now.strftime('%H:%M')}, {weekday}.")
        if topic == 'time':
            return

    ps_scripts = {
        'battery': '$b=Get-CimInstance Win32_Battery -EA SilentlyContinue; if($b){"$([int]$b.EstimatedChargeRemaining)%"}else{"батареи нет"}',
        'cpu': '$c=Get-CimInstance Win32_Processor -EA SilentlyContinue|Select-Object -First 1; $load=(Get-CimInstance Win32_PerfFormattedData_PerfOS_Processor -Filter "Name=\'_Total\'" -EA SilentlyContinue).PercentProcessorTime; if($null -eq $load){$load=$c.LoadPercentage}; "$load%"',
        'ram': '$o=Get-CimInstance Win32_OperatingSystem; $tg=[math]::Round($o.TotalVisibleMemorySize/1MB,1); $ug=[math]::Round(($o.TotalVisibleMemorySize-$o.FreePhysicalMemory)/1MB,1); "$ug из $tg гигабайт"',
        'disk': '$d=Get-CimInstance Win32_LogicalDisk -Filter "DeviceID=\'$env:SystemDrive\'"; $fg=[math]::Round($d.FreeSpace/1GB,0); $tg=[math]::Round($d.Size/1GB,0); "свободно $fg из $tg гигабайт"',
    }
    labels = {'battery': 'Питание', 'cpu': 'Процессор', 'ram': 'Память', 'disk': 'Диск'}
    order = ['battery', 'cpu', 'ram', 'disk'] if topic == 'all' else [topic]

    for t in order:
        ps = ps_scripts.get(t)
        if not ps:
            continue
        result = subprocess.run(
            ['powershell', '-NoProfile', '-Command', ps],
            capture_output=True, text=True, encoding='utf-8', timeout=10,
        )
        out = result.stdout.strip() if result.returncode == 0 else 'ошибка'
        _speak(f"{labels[t]}: {out}.")


def run_action(action):
    t = action['type']
    if t == 'exe':
        path = f"{CDIR}\\custom-commands\\{action['file']}"
        if action.get('wait'):
            subprocess.check_call([path])
        else:
            subprocess.Popen([path])
        if 'delay_ms' in action:
            time.sleep(action['delay_ms'] / 1000)
    elif t == 'play_sound':
        play(action['name'])
    elif t == 'system':
        op = action['op']
        if op == 'volume_mute':
            _set_mute(1)
        elif op == 'volume_unmute':
            _set_mute(0)
        elif op == 'exit':
            _shutdown()
    elif t == 'sleep':
        time.sleep(action['ms'] / 1000)
    elif t == 'shell':
        if action.get('wait'):
            subprocess.check_call(action['cmd'], shell=True)
        else:
            subprocess.Popen(action['cmd'], shell=True)
    elif t == 'url':
        browser = action.get('browser')
        if browser:
            exe = config.BROWSER_PATHS.get(browser)
            if exe and os.path.isfile(exe):
                subprocess.Popen([exe, action['href']])
                return
        webbrowser.open(action['href'])
    elif t == 'keys':
        import pyautogui
        if action.get('sequence'):
            keys = [k.strip() for k in action['sequence'].split('+') if k.strip()]
            pyautogui.hotkey(*keys)
        elif action.get('text'):
            pyautogui.typewrite(action['text'], interval=0.02)
    elif t == 'multi':
        for step in action['steps']:
            run_action(step)
    elif t == 'groq_codegen':
        do_codegen(action, _current_voice or '')
    elif t == 'ocr_screen':
        do_ocr(action, _current_voice or '')
    elif t == 'sysinfo':
        do_sysinfo(action, _current_voice or '')
    elif t == 'file_search':
        do_file_search(action, _current_voice or '')
    elif t == 'audio_devices_panel':
        do_audio_devices_panel(action, _current_voice or '')
    elif t == 'clipboard_read':
        do_clipboard_read(action, _current_voice or '')
    elif t == 'notes_add':
        do_notes_add(action, _current_voice or '')
    elif t == 'notes_open':
        do_notes_open(action, _current_voice or '')
    elif t == 'process_kill':
        do_process_kill(action, _current_voice or '')
    elif t == 'websearch':
        do_websearch(action, _current_voice or '')


_current_voice: str = ''


def execute_cmd(cmd: str, voice: str):
    global _current_voice
    spec = VA_CMD_LIST.get(cmd)
    if not spec:
        return
    action = spec['action']
    _current_voice = voice
    try:
        run_action(action)
    finally:
        _current_voice = ''
    confirm = spec.get('confirm_sound')
    if confirm is None and action['type'] == 'exe':
        confirm = 'ok'
    if confirm:
        play(confirm)


recorder = PvRecorder(device_index=config.MICROPHONE_INDEX, frame_length=512)
recorder.start()
print('Using device: %s' % recorder.selected_device)

print(f"Jarvis (v{config.VA_VER}) начал свою работу ...")
play("run")
time.sleep(0.5)


def heard_wake_word(text: str) -> bool:
    if not text:
        return False
    lowered = text.lower()
    return any(w in lowered for w in config.WAKE_WORDS)


while True:
    try:
        pcm = recorder.read()
        sp = struct.pack("h" * len(pcm), *pcm)

        if not wake_rec.AcceptWaveform(sp):
            continue

        result_text = json.loads(wake_rec.Result()).get("text", "")
        if not heard_wake_word(result_text):
            continue

        play("greet", True)
        print("Yes, sir.")
        kaldi_rec.Reset()
        listen_started = time.time()
        speech_ms = 0
        silence_ms = 0
        vad_buf = b""
        finalize = False
        cmd_buf = b"" if config.DENOISE_ENABLED else None

        while True:
            elapsed_ms = (time.time() - listen_started) * 1000
            if elapsed_ms >= config.COMMAND_MAX_LISTEN_MS:
                break

            pcm = recorder.read()
            sp = struct.pack("h" * len(pcm), *pcm)
            if cmd_buf is None:
                kaldi_rec.AcceptWaveform(sp)
            else:
                cmd_buf += sp

            vad_buf += sp
            while len(vad_buf) >= VAD_FRAME_BYTES:
                frame = vad_buf[:VAD_FRAME_BYTES]
                vad_buf = vad_buf[VAD_FRAME_BYTES:]
                if vad.is_speech(frame, samplerate):
                    speech_ms += VAD_FRAME_MS
                    silence_ms = 0
                else:
                    silence_ms += VAD_FRAME_MS

            if elapsed_ms < config.COMMAND_MIN_LISTEN_MS:
                continue
            if speech_ms >= config.COMMAND_MIN_SPEECH_MS and silence_ms >= config.COMMAND_END_SILENCE_MS:
                finalize = True
                break

        if finalize or speech_ms > 0:
            if cmd_buf is not None:
                kaldi_rec.AcceptWaveform(denoise_pcm(cmd_buf))
            text = json.loads(kaldi_rec.FinalResult()).get("text", "")
            if text:
                va_respond(text)

        wake_rec.Reset()

    except Exception as err:
        print(f"Unexpected {err=}, {type(err)=}")
        raise