"""Tray supervisor with authenticated local control and owned process cleanup."""
import argparse
import json
import logging
import multiprocessing as mp
import os
import secrets
import subprocess
import sys
import threading
import time
import webbrowser
from monitor_runtime.paths import data, resource, atomic_json

def worker(address, auth):
    from multiprocessing.connection import Client
    connection = Client(address, family='AF_PIPE', authkey=bytes.fromhex(auth))
    try:
        if connection.recv() != 'run':
            return 1
        from monitor_runtime.service import run
        return run(connection)
    finally:
        connection.close()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--worker', nargs=2, metavar=('PIPE', 'AUTH'))
    parser.add_argument('--stop', action='store_true')
    parser.add_argument('--self-test', action='store_true')
    parser.add_argument('--trust-model')
    parser.add_argument('--yes', action='store_true')
    args = parser.parse_args()
    if os.name != 'nt':
        raise RuntimeError('Monitor desktop requires Windows')
    data('.runtime').mkdir(parents=True, exist_ok=True)
    data('log').mkdir(parents=True, exist_ok=True)
    logging.basicConfig(filename=str(data('log/desktop.log')), level=logging.INFO,
                        format='%(asctime)s %(levelname)s %(message)s')
    if args.self_test:
        from monitor_runtime.diagnostics import run
        return run()
    if args.worker:
        return worker(*args.worker)
    if args.trust_model:
        return trust_model(args.trust_model, args.yes)
    from app.services.lifecycle import ServiceLeaderLock
    # Separate from the service lock; held for the entire supervisor lifetime.
    lock = ServiceLeaderLock(data('.runtime/desktop.lock'))
    if not lock.acquire():
        return control_existing('stop' if args.stop else 'open')
    if args.stop:
        lock.release()
        return 0
    try:
        return Supervisor().run()
    finally:
        lock.release()

def control_existing(command):
    from multiprocessing.connection import Client
    deadline = time.monotonic() + 10
    while True:
        try:
            state = json.loads(data('.runtime/desktop.json').read_text(encoding='utf-8'))
            with Client(state['pipe'], family='AF_PIPE', authkey=bytes.fromhex(state['auth'])) as connection:
                connection.send(command)
                response = connection.recv()
            break
        except (OSError, EOFError, json.JSONDecodeError):
            if time.monotonic() >= deadline:
                raise RuntimeError('已有实例的控制接口不可用，请查看日志')
            time.sleep(.2)
    if command == 'stop':
        deadline = time.monotonic() + 40
        from app.services.lifecycle import ServiceLeaderLock
        while time.monotonic() < deadline:
            lock = ServiceLeaderLock(data('.runtime/desktop.lock'))
            if lock.acquire():
                lock.release()
                return 0
            time.sleep(.3)
        raise RuntimeError('Monitor 未能在 40 秒内退出，请查看日志')
    return 0

def trust_model(filename, confirmed):
    from pathlib import Path
    from app.utils.ModelTrust import sha256_file
    from monitor_runtime.windows import message
    model = Path(filename).resolve()
    if model.suffix.lower() != '.pt' or not model.is_file():
        raise ValueError('请选择存在的 .pt 模型文件')
    digest = sha256_file(model)
    if not confirmed:
        message('模型 SHA-256: ' + digest + '\n确认来源可信后使用 --yes 重新执行。')
        return 2
    path = data('.trusted-models.json')
    doc = json.loads(path.read_text(encoding='utf-8')) if path.exists() else {'trusted_sha256': []}
    doc['trusted_sha256'] = sorted(set(doc['trusted_sha256']) | {digest})
    atomic_json(path, doc)
    message('已登记可信模型 SHA-256：' + digest)
    return 0

class Supervisor:
    def __init__(self):
        self.stop = threading.Event()
        self.restart = threading.Event()
        self.url = ''
        self.status = '正在启动'
        self.icon = None
        self.child = None
        self.pipe = None
        self.job = None

    def open_web(self, *_):
        if self.url:
            webbrowser.open(self.url)

    def exit(self, *_):
        self.stop.set()

    def run(self):
        import pystray
        from PIL import Image
        from multiprocessing.connection import Listener
        auth = secrets.token_bytes(32)
        address = r'\\.\pipe\MonitorControl-' + secrets.token_hex(16)
        self.control = Listener(address, family='AF_PIPE', authkey=auth)
        atomic_json(data('.runtime/desktop.json'), {'pipe': address, 'auth': auth.hex()})
        threading.Thread(target=self.control_loop, daemon=True).start()
        menu = pystray.Menu(
            pystray.MenuItem('打开管理网页', self.open_web, default=True),
            pystray.MenuItem(lambda _: self.status, lambda *_: None, enabled=False),
            pystray.MenuItem('查看状态', lambda *_: self.show_status()),
            pystray.MenuItem('打开日志目录', lambda *_: os.startfile(str(data('log')))),
            pystray.MenuItem('重启服务', lambda *_: self.restart.set()),
            pystray.MenuItem('退出', self.exit))
        self.icon = pystray.Icon('Monitor', Image.open(resource('static/images/logo.png')), 'Monitor', menu)
        threading.Thread(target=self.supervise, daemon=True).start()
        try:
            self.icon.run()
        finally:
            self.stop.set()
            self.control.close()
        return 0

    def show_status(self):
        from monitor_runtime.windows import message
        message(self.status + '\n' + self.url.split('#')[0])

    def control_loop(self):
        while not self.stop.is_set():
            try:
                with self.control.accept() as conn:
                    command = conn.recv()
                    if command == 'stop':
                        self.stop.set()
                    elif command == 'open':
                        self.open_web()
                    conn.send('ok')
            except (OSError, EOFError):
                return

    def start_child(self):
        from multiprocessing.connection import Listener
        from monitor_runtime.windows import Job
        auth = secrets.token_bytes(32)
        address = r'\\.\pipe\MonitorWorker-' + secrets.token_hex(16)
        listener = Listener(address, family='AF_PIPE', authkey=auth)
        args = [sys.executable]
        if not getattr(sys, 'frozen', False):
            args.append(str(resource('monitor_entry.py')))
        args += ['--worker', address, auth.hex()]
        env = os.environ.copy()
        env['MONITOR_DESKTOP'] = '1'
        env['MONITOR_SERVICE_MODE'] = 'disabled'
        env['MONITOR_BOOTSTRAP_SERVICES'] = 'false'
        env['MONITOR_DEBUG'] = 'false'
        env['MONITOR_DATA_DIR'] = str(data(''))
        self.job = Job()
        log = open(data('log/service.log'), 'ab', buffering=0)
        try:
            self.child = subprocess.Popen(args, cwd=str(data('')), env=env, stdin=subprocess.DEVNULL,
                                          stdout=log, stderr=log, creationflags=subprocess.CREATE_NO_WINDOW)
            self.job.assign(self.child)
        finally:
            log.close()
        # Child waits for run; Job assignment therefore precedes any child spawn.
        accepted = []
        def accept():
            try:
                accepted.append(listener.accept())
            except OSError:
                pass
        thread = threading.Thread(target=accept, daemon=True)
        thread.start()
        deadline = time.monotonic() + 30
        while not accepted:
            if self.child.poll() is not None or time.monotonic() > deadline or self.stop.wait(.1):
                listener.close()
                raise RuntimeError('服务进程未能建立控制连接，请查看 service.log')
        self.pipe = accepted[0]
        listener.close()
        self.pipe.send('run')

    def stop_child(self):
        if self.child and self.child.poll() is None:
            try:
                self.pipe.send('stop')
            except (AttributeError, EOFError, OSError):
                pass
            try:
                self.child.wait(timeout=25)
            except subprocess.TimeoutExpired:
                logging.error('Graceful shutdown timed out; closing owned job')
        if self.pipe:
            self.pipe.close()
            self.pipe = None
        if self.job:
            self.job.close()
            self.job = None
        self.child = None

    def supervise(self):
        try:
            while not self.stop.is_set():
                self.restart.clear()
                self.status = '正在启动'
                self.start_child()
                deadline = time.monotonic() + 180
                last_heartbeat = time.monotonic()
                opened = False
                auto_restart = False
                try:
                    while not self.stop.is_set() and not self.restart.is_set():
                        if self.pipe.poll(.3):
                            message = self.pipe.recv()
                            last_heartbeat = time.monotonic()
                            phase = message.get('phase')
                            if phase == 'failed':
                                raise RuntimeError(message.get('reason', '服务启动失败'))
                            self.status = {'ready': '服务运行中', 'setup': '等待首次初始化',
                                           'license': '等待有效授权'}.get(phase, phase)
                            if message.get('url'):
                                self.url = message['url']
                                if not opened:
                                    self.open_web()
                                    opened = True
                            self.icon.update_menu()
                        code = self.child.poll()
                        if code is not None:
                            if code == 20:
                                auto_restart = True
                                break
                            raise RuntimeError('服务进程退出，代码 %s' % code)
                        if not opened and time.monotonic() > deadline:
                            raise RuntimeError('启动超过 180 秒，请查看日志')
                        if opened and time.monotonic() - last_heartbeat > 35:
                            raise RuntimeError('服务健康状态超时')
                except EOFError:
                    code = self.child.wait(timeout=30)
                    if code == 20:
                        auto_restart = True
                    else:
                        raise RuntimeError('服务控制连接断开，代码 %s' % code)
                finally:
                    self.stop_child()
                if not auto_restart and not self.restart.is_set():
                    break
        except Exception as exc:
            logging.exception('Supervisor failure')
            self.status = '启动失败：' + str(exc)
            self.url = ''
            self.icon.update_menu()
            from monitor_runtime.windows import message
            message(self.status + '\n日志：' + str(data('log')), True)
            self.stop_child()
            # Keep tray available; retries are explicit, never an infinite crash loop.
            while not self.stop.wait(.3):
                if self.restart.is_set():
                    self.supervise()
                    return
        finally:
            self.stop_child()
            if self.stop.is_set():
                self.icon.stop()