J.A.R.V.I.S-py/extensions.py

"""Extension handlers — new commands ported from the Rust fork.

Each function takes (action, voice) and returns nothing; it speaks via the
`speak` callback passed in from main.py. We keep the speech callback as
a module global to avoid passing it through every call site.

Usage from main.py:
    import extensions
    extensions.set_speak(_speak)
    # then dispatch:
    if t == 'lock_workstation': extensions.do_lock(action, voice)
"""

import os
import math
import random
import subprocess
import time
import uuid as uuidlib
import re
import urllib.request

import memory_store
import profiles_store

_speak_fn = print
_set_clipboard_fn = None


def set_speak(fn):
    """Register the main.py _speak function as our voice output."""
    global _speak_fn
    _speak_fn = fn


def set_clipboard_fn(fn):
    """Register the main.py _set_clipboard helper."""
    global _set_clipboard_fn
    _set_clipboard_fn = fn


def _speak(text):
    try:
        _speak_fn(text)
    except Exception as exc:
        print(f"[ext] speak failed: {exc}")


def _clip(text):
    if _set_clipboard_fn:
        try:
            _set_clipboard_fn(text)
        except Exception as exc:
            print(f"[ext] clipboard set failed: {exc}")


def _ps(cmd, timeout=15):
    """Run a PowerShell command, return stripped stdout."""
    res = subprocess.run(
        ['powershell', '-NoProfile', '-ExecutionPolicy', 'Bypass', '-Command', cmd],
        capture_output=True, text=True, encoding='utf-8', timeout=timeout,
    )
    return (res.stdout or '').strip()


# ── lock_workstation ───────────────────────────────────────────────────────

def do_lock_workstation(action, voice):
    """Lock the workstation via LockWorkStation() in user32.dll."""
    subprocess.Popen(['rundll32.exe', 'user32.dll,LockWorkStation'])
    _speak("Компьютер заблокирован.")


# ── screenshot ─────────────────────────────────────────────────────────────

_SHOT_TO_CLIP_PS = (
    "Add-Type -AssemblyName System.Windows.Forms; "
    "Add-Type -AssemblyName System.Drawing; "
    "$b = [System.Windows.Forms.Screen]::PrimaryScreen.Bounds; "
    "$bmp = New-Object System.Drawing.Bitmap $b.Width, $b.Height; "
    "$g = [System.Drawing.Graphics]::FromImage($bmp); "
    "$g.CopyFromScreen($b.Location, [System.Drawing.Point]::Empty, $b.Size); "
    "[System.Windows.Forms.Clipboard]::SetImage($bmp); "
    "$g.Dispose(); $bmp.Dispose()"
)


def do_screenshot(action, voice):
    mode = action.get('mode', 'clipboard')
    if mode == 'clipboard':
        try:
            subprocess.run(
                ['powershell', '-NoProfile', '-Command', _SHOT_TO_CLIP_PS],
                capture_output=True, timeout=10, check=True,
            )
            _speak("Скриншот в буфере.")
        except Exception as exc:
            print(f"[ext] screenshot clipboard failed: {exc}")
            _speak("Не получилось сделать скриншот.")
    elif mode == 'file':
        userprofile = os.environ.get('USERPROFILE', 'C:\\Users\\Public')
        screens_dir = os.path.join(userprofile, 'Pictures', 'Screenshots')
        os.makedirs(screens_dir, exist_ok=True)
        ts = time.strftime('%Y%m%d-%H%M%S')
        path = os.path.join(screens_dir, f'jarvis-{ts}.png')
        ps = (
            "Add-Type -AssemblyName System.Windows.Forms; "
            "Add-Type -AssemblyName System.Drawing; "
            "$b = [System.Windows.Forms.Screen]::PrimaryScreen.Bounds; "
            "$bmp = New-Object System.Drawing.Bitmap $b.Width, $b.Height; "
            "$g = [System.Drawing.Graphics]::FromImage($bmp); "
            "$g.CopyFromScreen($b.Location, [System.Drawing.Point]::Empty, $b.Size); "
            f"$bmp.Save('{path}', [System.Drawing.Imaging.ImageFormat]::Png); "
            "$g.Dispose(); $bmp.Dispose()"
        )
        try:
            subprocess.run(['powershell', '-NoProfile', '-Command', ps],
                           capture_output=True, timeout=10, check=True)
            _speak("Скриншот сохранён в папку Скриншоты.")
        except Exception as exc:
            print(f"[ext] screenshot file failed: {exc}")
            _speak("Не получилось сохранить скриншот.")


# ── disk space ─────────────────────────────────────────────────────────────

def do_disk_free(action, voice):
    """Reports free GB on the requested drive (defaults to C)."""
    letter = (action.get('letter') or 'C').upper()
    # Pull letter from voice if user named one
    m = re.search(r'диск[еа]?\s*([A-ZА-Я])', (voice or '').upper())
    if m:
        letter = m.group(1)
        # cyrillic Е/А fallback to C
        if letter not in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ':
            letter = 'C'

    out = _ps(
        f"$d = Get-PSDrive {letter} -ErrorAction SilentlyContinue; "
        "if ($d) {{ '{{0}}|{{1}}' -f "
        "[math]::Round($d.Free/1GB,1), "
        "[math]::Round(($d.Free + $d.Used)/1GB,0) }} else {{ 'NONE' }}".format(
            letter=letter
        )
    )
    if out == 'NONE' or '|' not in out:
        _speak(f"Диск {letter} не найден.")
        return
    free, total = out.split('|', 1)
    try:
        pct = int(round(float(free) / float(total) * 100))
        _speak(f"На диске {letter} свободно {free} гигабайт из {total}, это {pct} процентов.")
    except (ValueError, ZeroDivisionError):
        _speak(f"На диске {letter} свободно {free} гигабайт.")


def do_disk_list(action, voice):
    out = _ps(
        "Get-PSDrive -PSProvider FileSystem | "
        "Where-Object { $_.Used -ne $null -or $_.Free -ne $null } | "
        "ForEach-Object { '{0}:{1}' -f $_.Name, [math]::Round($_.Free/1GB,0) }"
    )
    if not out:
        _speak("Дисков не нашёл.")
        return
    parts = []
    for tok in out.split():
        m = re.match(r'([A-Z]):(\d+)', tok)
        if m:
            parts.append(f"{m.group(1)} {m.group(2)} гигабайт")
    if not parts:
        _speak("Не понял ответ системы.")
        return
    _speak("Свободно: " + ", ".join(parts) + ".")


# ── random generators ──────────────────────────────────────────────────────

def do_coin(action, voice):
    if random.randint(1, 2) == 1:
        _speak("Орёл!")
    else:
        _speak("Решка!")


_PASS_POOL = (
    'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*-_=+'
)


def do_password_gen(action, voice):
    m = re.search(r'(\d+)', voice or '')
    n = int(m.group(1)) if m else 16
    n = max(6, min(64, n))
    password = ''.join(random.choice(_PASS_POOL) for _ in range(n))
    _clip(password)
    _speak(f"Пароль на {n} символов в буфере.")


def do_uuid(action, voice):
    u = str(uuidlib.uuid4())
    _clip(u)
    _speak(f"UUID в буфере, заканчивается на {u[-4:]}.")


# ── SSL cert check ─────────────────────────────────────────────────────────

def do_ssl_check(action, voice):
    domain = action.get('domain') or _extract_domain(voice)
    if not domain:
        _speak("Какой домен проверить?")
        return
    domain = domain.replace('https://', '').replace('http://', '').split('/')[0]
    ps = f"""
$ErrorActionPreference = 'Stop'
try {{
  $tcp = New-Object Net.Sockets.TcpClient('{domain}', 443)
  $ssl = New-Object Net.Security.SslStream($tcp.GetStream(), $false, {{ $true }})
  $ssl.AuthenticateAsClient('{domain}')
  $c = New-Object Security.Cryptography.X509Certificates.X509Certificate2($ssl.RemoteCertificate)
  $days = ($c.NotAfter - (Get-Date)).Days
  Write-Output ('{{0}}|{{1}}' -f $days, $c.NotAfter.ToString('yyyy-MM-dd'))
  $ssl.Close(); $tcp.Close()
}} catch {{ Write-Output ('ERR|' + $_.Exception.Message) }}
"""
    out = _ps(ps)
    if out.startswith('ERR|'):
        _speak(f"Не получилось: {out[4:80]}")
        return
    if '|' not in out:
        _speak("Не распарсил ответ.")
        return
    days, date = out.split('|', 1)
    try:
        days_n = int(days)
    except ValueError:
        _speak("Не понял ответ.")
        return
    prefix = ''
    if days_n < 0:
        prefix = "Сертификат уже просрочен! "
    elif days_n < 7:
        prefix = "СКОРО! "
    elif days_n < 30:
        prefix = "Скоро. "
    _speak(f"{prefix}Сертификат {domain} истекает через {days} дней ({date}).")


def _extract_domain(voice):
    if not voice:
        return ''
    # crude: pick last token that looks like a domain (has a dot)
    for tok in voice.split():
        if '.' in tok and len(tok) > 3:
            return tok
    return ''


# ── WiFi SSID ──────────────────────────────────────────────────────────────

def do_wifi_ssid(action, voice):
    res = subprocess.run(['netsh', 'wlan', 'show', 'interfaces'],
                         capture_output=True, text=True, encoding='cp866', timeout=5)
    out = res.stdout or ''
    # Try ru + en field names
    m = re.search(r'^\s*SSID\s+:\s+(.+?)$', out, re.MULTILINE)
    if not m:
        m = re.search(r'Имя сети\s+:\s+(.+?)$', out, re.MULTILINE)
    if m:
        ssid = m.group(1).strip()
        if ssid in ('(none)', '(нет)'):
            _speak("Wi-Fi не подключен.")
            return
        _speak(f"Сеть: {ssid}.")
    else:
        _speak("Wi-Fi не подключен или адаптер выключен.")


def do_public_ip(action, voice):
    try:
        with urllib.request.urlopen('https://api.ipify.org', timeout=5) as r:
            ip = r.read().decode().strip()
        if ip:
            _speak(f"Внешний адрес: {ip}.")
        else:
            _speak("Сервис вернул пустой ответ.")
    except Exception as exc:
        print(f"[ext] public_ip failed: {exc}")
        _speak("Не получилось узнать внешний IP.")


# ── unit conversion ────────────────────────────────────────────────────────

def _grab_number(voice):
    m = re.search(r'(-?[\d.,]+)', voice or '')
    if not m:
        return None
    try:
        return float(m.group(1).replace(',', '.'))
    except ValueError:
        return None


def _ru_plural(n, one, few, many):
    """Russian-grammar pluralisation: 1 фут, 2-4 фута, 5+ футов."""
    nn = abs(int(n))
    last_two = nn % 100
    last = nn % 10
    if 11 <= last_two <= 19:
        return many
    if last == 1:
        return one
    if 2 <= last <= 4:
        return few
    return many


def do_unit_convert(action, voice):
    mode = action.get('mode', 'length')
    n = _grab_number(voice)
    if n is None:
        _speak("Не понял число.")
        return
    low = (voice or '').lower()
    if mode == 'length':
        if 'фут' in low and 'в метр' not in low:
            result = n * 3.28084
            unit = _ru_plural(result, 'фут', 'фута', 'футов')
        elif 'метр' in low and 'в фут' not in low:
            result = n / 3.28084
            unit = 'метров'
        elif 'мил' in low and 'в км' not in low:
            result = n * 0.621371
            unit = _ru_plural(result, 'миля', 'мили', 'миль')
        elif 'км' in low or 'километр' in low:
            result = n * 1.60934
            unit = 'километров'
        else:
            _speak("Не понял в чём перевести.")
            return
        _speak(f"{round(result, 2):g} {unit}.")
    elif mode == 'weight':
        if 'фунт' in low and 'в кг' not in low:
            result = n * 2.20462
            unit = _ru_plural(result, 'фунт', 'фунта', 'фунтов')
        else:
            result = n / 2.20462
            unit = 'килограмм'
        _speak(f"{round(result, 2):g} {unit}.")
    elif mode == 'temperature':
        if 'фаренгейт' in low or 'fahrenheit' in low:
            result = n * 9 / 5 + 32
            unit = 'по Фаренгейту'
        else:
            result = (n - 32) * 5 / 9
            unit = 'по Цельсию'
        _speak(f"{round(result, 1):g} градусов {unit}.")
    elif mode == 'speed':
        if 'мил' in low and 'в км' not in low:
            result = n * 0.621371
            unit = 'миль в час'
        else:
            result = n * 1.60934
            unit = 'километров в час'
        _speak(f"{int(round(result))} {unit}.")
    else:
        _speak(f"Неизвестный режим конверсии: {mode}.")


# ── date math ──────────────────────────────────────────────────────────────

_MONTH_STEMS = {
    'январ': 1, 'феврал': 2, 'март': 3, 'апрел': 4,
    'мая': 5, 'май': 5, 'июн': 6, 'июл': 7,
    'август': 8, 'сентябр': 9, 'октябр': 10, 'ноябр': 11, 'декабр': 12,
}

_DAYS_RU = [
    'понедельник', 'вторник', 'среда', 'четверг',
    'пятница', 'суббота', 'воскресенье',
]


def do_date_math(action, voice):
    import datetime as dt
    mode = action.get('mode', 'days_until')
    low = (voice or '').lower()
    today = dt.date.today()

    if mode == 'day_of_week':
        # weekday() returns 0=Monday..6=Sunday
        _speak(f"Сегодня {_DAYS_RU[today.weekday()]}.")
        return

    # days_until
    target = None
    if 'нов' in low and 'год' in low:
        ny = today.year if (today.month == 1 and today.day == 1) else today.year + 1
        target = dt.date(ny, 1, 1)
    elif 'рожд' in low:
        target = dt.date(today.year, 1, 7)
    elif '8 март' in low or 'восьмого март' in low:
        target = dt.date(today.year, 3, 8)
    elif '9 ма' in low or 'девятое ма' in low:
        target = dt.date(today.year, 5, 9)
    else:
        m = re.search(r'до\s+(\d+)', low)
        if m:
            day = int(m.group(1))
            for stem, month in _MONTH_STEMS.items():
                if stem in low:
                    try:
                        target = dt.date(today.year, month, day)
                    except ValueError:
                        target = None
                    break

    if target is None:
        _speak("Не понял дату.")
        return
    if target < today:
        target = target.replace(year=target.year + 1)
    diff = (target - today).days
    if diff == 0:
        _speak("Сегодня.")
        return
    unit = _ru_plural(diff, 'день', 'дня', 'дней')
    _speak(f"{diff} {unit}.")


# ── echo / self-check ──────────────────────────────────────────────────────

def do_echo_repeat(action, voice):
    body = (voice or '').strip()
    for trig in ('повтори за мной', 'скажи как я', 'эхо', 'repeat after me', 'echo'):
        idx = body.lower().find(trig)
        if idx == 0:
            body = body[len(trig):].strip(' ,.:')
            break
    if not body:
        _speak("Что повторить?")
        return
    _speak(body)


def do_self_ping(action, voice):
    options = ['Да, сэр.', 'Здесь, сэр.', 'Слушаю.', 'К вашим услугам.', 'На месте.']
    _speak(random.choice(options))


# ── memory ──────────────────────────────────────────────────────────────────

_MEMORY_REMEMBER_TRIGGERS = (
    'запомни обо мне что', 'запомни обо мне', 'запомни что',
    'помни что', 'запомни это', 'запомни про меня', 'запомни',
)


def do_memory_remember(action, voice):
    body = (voice or '').strip()
    low = body.lower()
    for trig in _MEMORY_REMEMBER_TRIGGERS:
        if low.startswith(trig):
            body = body[len(trig):].strip(' ,.:')
            break

    if not body:
        _speak("Что именно запомнить?")
        return

    if '=' in body:
        key, value = body.split('=', 1)
        key, value = key.strip(), value.strip()
    else:
        words = body.split()
        key = ' '.join(words[:6])
        value = body

    if not key:
        _speak("Не понял ключ.")
        return

    memory_store.remember(key, value)
    _speak("Запомнил.")


def do_memory_recall(action, voice):
    body = (voice or '').lower()
    query = ''
    for trig in ('что ты помнишь о', 'что ты помнишь про', 'что ты знаешь обо мне',
                 'что помнишь про', 'вспомни'):
        idx = body.find(trig)
        if idx >= 0:
            query = body[idx + len(trig):].strip(' ,.:')
            break

    if not query:
        recs = memory_store.all_facts()
        if not recs:
            _speak("Ничего пока не помню.")
            return
        line = "Я помню: "
        for r in recs[:3]:
            line += f"{r['key']} — {r['value']}. "
        _speak(line)
        return

    hits = memory_store.search(query, 3)
    if not hits:
        _speak(f"Ничего не нашёл про {query}.")
        return
    if len(hits) == 1:
        _speak(hits[0]['value'])
    else:
        _speak('. '.join(f"{h['key']}: {h['value']}" for h in hits))


def do_memory_forget(action, voice):
    body = (voice or '').lower()
    for trig in ('забудь про', 'забудь что', 'забудь обо мне'):
        idx = body.find(trig)
        if idx >= 0:
            key = body[idx + len(trig):].strip(' ,.:')
            if memory_store.forget(key):
                _speak(f"Забыл про {key}.")
                return
            hits = memory_store.search(key, 1)
            if hits:
                memory_store.forget(hits[0]['key'])
                _speak(f"Забыл про {hits[0]['key']}.")
                return
            _speak(f"Я и не помнил про {key}.")
            return
    _speak("Что забыть?")


def do_memory_list(action, voice):
    recs = memory_store.all_facts()
    if not recs:
        _speak("Я пока ничего о вас не запомнил.")
        return
    sample = min(len(recs), 5)
    line = f"Помню {len(recs)} фактов. Первые: "
    line += ", ".join(r['key'] for r in recs[:sample]) + "."
    _speak(line)


# ── profile switch ──────────────────────────────────────────────────────────

def do_profile_set(action, voice):
    target = action.get('target') or _profile_target_from_voice(voice)
    if not target:
        _speak("Не понял какой режим.")
        return
    try:
        p = profiles_store.set_active(target)
    except ValueError:
        _speak(f"Профиль {target} не найден.")
        return
    greeting = p.get('greeting', '') or f"Режим {target}."
    _speak(greeting)


def _profile_target_from_voice(voice):
    low = (voice or '').lower()
    if 'работ' in low: return 'work'
    if 'игр' in low: return 'game'
    if 'сон' in low or 'спать' in low or 'тиш' in low: return 'sleep'
    if 'руль' in low or 'вожд' in low or 'поехали' in low: return 'driving'
    if 'обычн' in low or 'умолчан' in low or 'сброс' in low or 'верн' in low: return 'default'
    return ''


def do_profile_what(action, voice):
    p = profiles_store.active()
    name = p.get('name', 'default')
    desc = p.get('description', '')
    msg = f"Сейчас {name}."
    if desc:
        msg += " " + desc
    _speak(msg)


# ── interesting fact (LLM-dependent) ───────────────────────────────────────

def do_interesting_fact(action, voice):
    """Asks the configured Groq LLM for a fun fact. Requires GROQ_TOKEN."""
    import config as cfg
    if not getattr(cfg, 'GROQ_TOKEN', None):
        _speak("LLM не настроен.")
        return

    topic = ''
    low = (voice or '').lower()
    for trig in ('расскажи интересный факт про', 'расскажи интересный факт о',
                 'интересный факт про', 'интересный факт о',
                 'расскажи интересный факт', 'интересный факт',
                 'расскажи факт', 'удиви меня'):
        idx = low.find(trig)
        if idx >= 0:
            topic = (voice or '')[idx + len(trig):].strip(' ,.:').strip()
            break

    if topic:
        prompt = f"Расскажи один интересный факт про {topic}. На русском, 1-2 предложения. Без вступлений."
    else:
        prompt = "Расскажи один реально неочевидный научно-проверенный факт. На русском, 1-2 предложения. Без вступлений типа 'знали ли вы'."

    try:
        from openai import OpenAI
    except ImportError:
        _speak("OpenAI клиент не установлен.")
        return

    client = OpenAI(
        api_key=cfg.GROQ_TOKEN,
        base_url=getattr(cfg, 'GROQ_BASE_URL', 'https://api.groq.com/openai/v1'),
    )
    try:
        resp = client.chat.completions.create(
            model=getattr(cfg, 'GROQ_MODEL', 'llama-3.3-70b-versatile'),
            messages=[
                {'role': 'system', 'content': 'Ты любопытный собеседник. Цепляющие факты, без воды.'},
                {'role': 'user', 'content': prompt},
            ],
            max_tokens=200, temperature=0.8, timeout=30,
        )
        reply = resp.choices[0].message.content.strip()
        if reply:
            _speak(reply)
        else:
            _speak("Сегодня нет интересных фактов.")
    except Exception as exc:
        print(f"[ext] fact LLM failed: {exc}")
        _speak("Не получилось получить факт.")