#!/usr/bin/env python3
"""
companion.py — 웹 코드 에디터용 로컬 Python 실행 컴패니언

사용법:
    pip install websockets
    python companion.py

웹 에디터(브라우저)가 ws://localhost:9000 에 연결해서
코드 실행 요청을 보내면 이 프로세스가 로컬에서 실행합니다.
pygame 등 GUI 라이브러리도 정상 동작합니다.
"""

import asyncio
import io
import json
import locale
import os
import sys
import tempfile
import threading

# pythonw.exe에서는 sys.stdout/stderr가 None이므로 print() 크래시 방지
if sys.stdout is None:
    sys.stdout = io.TextIOWrapper(io.BytesIO(), encoding='utf-8', errors='replace')
if sys.stderr is None:
    sys.stderr = io.TextIOWrapper(io.BytesIO(), encoding='utf-8', errors='replace')

try:
    import websockets
except ImportError:
    print('오류: websockets 패키지가 없습니다.')
    print('설치: pip install websockets')
    sys.exit(1)

HOST = '127.0.0.1'
PORT = 9000
IS_WINDOWS = sys.platform == 'win32'

# 자동시작(.vbs)으로 실행되면 프로세스 cwd가 System32로 잡히는 경우가 있어
# (wscript.exe 기본 작업 폴더 상속), 터미널이 엉뚱한 위치에서 열리는 걸 막기
# 위해 항상 홈 디렉터리에서 시작하도록 고정한다.
try:
    os.chdir(os.path.expanduser('~'))
except Exception:
    pass


def _refresh_env():
    """시작 후 변경된 환경 변수(PATH 등)를 레지스트리에서 실시간으로 반영한다."""
    if not IS_WINDOWS:
        return
    try:
        import winreg
        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                            r'SYSTEM\CurrentControlSet\Control\Session Manager\Environment',
                            0, winreg.KEY_READ) as key:
            sys_path = winreg.QueryValueEx(key, 'Path')[0]
        with winreg.OpenKey(winreg.HKEY_CURRENT_USER,
                            r'Environment',
                            0, winreg.KEY_READ) as key:
            user_path = winreg.QueryValueEx(key, 'Path')[0]
        expanded = os.path.expandvars(sys_path + ';' + user_path)
        existing = set(p.strip().lower() for p in os.environ.get('Path', '').split(';') if p.strip())
        new_entries = [p for p in expanded.split(';') if p.strip() and p.strip().lower() not in existing]
        os.environ['Path'] = os.environ.get('Path', '') + ';' + ';'.join(new_entries)
    except Exception:
        pass


async def _safe_send(ws, data: str):
    try:
        await ws.send(data)
    except Exception:
        pass


class ShellSession:
    """shellId별 독립 PTY 셸 세션."""

    def __init__(self, shell_id: str):
        self.shell_id = shell_id
        self.proc = None
        self.master_fd = None
        self._shell_task = None
        self._win_pty = None

    async def kill(self):
        if IS_WINDOWS and self._win_pty:
            try:
                if hasattr(self._win_pty, 'close'):
                    self._win_pty.close()
            except Exception:
                pass
            self._win_pty = None
        if self._shell_task:
            self._shell_task.cancel()
            try:
                await self._shell_task
            except (asyncio.CancelledError, Exception):
                pass
            self._shell_task = None
        if self.master_fd is not None:
            try:
                asyncio.get_running_loop().remove_reader(self.master_fd)
            except Exception:
                pass
            try:
                os.close(self.master_fd)
            except Exception:
                pass
            self.master_fd = None
        if self.proc and self.proc.returncode is None:
            try:
                self.proc.kill()
                await self.proc.wait()
            except Exception:
                pass
        self.proc = None

    def _ready_msg(self):
        return json.dumps({
            'type': 'shell_ready',
            'shellId': self.shell_id,
            'platform': 'windows' if IS_WINDOWS else 'unix',
            'home': os.path.expanduser('~'),
        })

    async def start(self, ws, rows: int = 24, cols: int = 80):
        if IS_WINDOWS:
            await self._start_windows(ws, rows, cols)
        else:
            await self._start_unix(ws, rows, cols)

    async def _start_unix(self, ws, rows, cols):
        import pty, fcntl, termios, struct
        master_fd, slave_fd = pty.openpty()
        winsize = struct.pack('HHHH', rows, cols, 0, 0)
        fcntl.ioctl(slave_fd, termios.TIOCSWINSZ, winsize)
        fl = fcntl.fcntl(master_fd, fcntl.F_GETFL)
        fcntl.fcntl(master_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
        self.proc = await asyncio.create_subprocess_exec(
            '/bin/bash', '--login',
            stdin=slave_fd, stdout=slave_fd, stderr=slave_fd,
            close_fds=True,
            env={**os.environ, 'TERM': 'xterm-256color'},
        )
        os.close(slave_fd)
        self.master_fd = master_fd
        await ws.send(self._ready_msg())
        self._shell_task = asyncio.create_task(self._stream_unix(ws))

    async def _start_windows(self, ws, rows, cols):
        # 1) pywinpty 2.x
        try:
            from pywinpty import Pty
            self._win_pty = Pty(cols, rows)
            self._win_pty.spawn('cmd.exe')
            await ws.send(self._ready_msg())
            self._shell_task = asyncio.create_task(self._stream_winpty(ws))
            print(f'[셸:{self.shell_id}] pywinpty 2.x', flush=True)
            return
        except ImportError:
            pass
        except Exception as e:
            print(f'[셸:{self.shell_id}] pywinpty 2.x 실패: {e}', flush=True)

        # 2) pywinpty 0.x
        try:
            from winpty import PtyProcess
            self._win_pty = PtyProcess.spawn(['cmd.exe'], dimensions=(rows, cols))
            await ws.send(self._ready_msg())
            self._shell_task = asyncio.create_task(self._stream_winpty(ws))
            print(f'[셸:{self.shell_id}] pywinpty 0.x', flush=True)
            return
        except ImportError:
            pass

        # 3) pipe fallback
        print(f'[셸:{self.shell_id}] pipe 모드 (pip install pywinpty 권장)', flush=True)
        self.proc = await asyncio.create_subprocess_exec(
            'cmd.exe', '/K', 'chcp 65001 > nul',
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            env={**os.environ},
        )
        await ws.send(self._ready_msg())
        self._shell_task = asyncio.create_task(self._stream_pipe(ws))

    async def _stream_unix(self, ws):
        loop = asyncio.get_running_loop()
        done = asyncio.Event()

        def on_readable():
            fd = self.master_fd
            if fd is None:
                done.set()
                return
            try:
                data = os.read(fd, 4096)
                loop.create_task(_safe_send(ws, json.dumps({
                    'type': 'shell_output',
                    'shellId': self.shell_id,
                    'data': data.decode('utf-8', errors='replace'),
                })))
            except BlockingIOError:
                pass
            except OSError:
                try:
                    loop.remove_reader(fd)
                except Exception:
                    pass
                done.set()

        loop.add_reader(self.master_fd, on_readable)
        print(f'[셸스트림:{self.shell_id}] 시작 fd={self.master_fd}', flush=True)
        try:
            await done.wait()
        except asyncio.CancelledError:
            fd = self.master_fd
            if fd is not None:
                try:
                    loop.remove_reader(fd)
                except Exception:
                    pass
            return

        if self.proc:
            try:
                await self.proc.wait()
            except Exception:
                pass
        try:
            await ws.send(json.dumps({
                'type': 'shell_exit',
                'shellId': self.shell_id,
                'code': self.proc.returncode if self.proc else -1,
            }))
        except Exception:
            pass

    async def _stream_winpty(self, ws):
        loop = asyncio.get_running_loop()
        try:
            while self._win_pty:
                is_alive = False
                try:
                    if hasattr(self._win_pty, 'is_alive'):
                        is_alive = self._win_pty.is_alive()
                    elif hasattr(self._win_pty, 'isalive'):
                        is_alive = self._win_pty.isalive()
                except Exception:
                    break
                if not is_alive:
                    break
                try:
                    data = await loop.run_in_executor(None, self._win_pty.read, 4096)
                    if data:
                        await _safe_send(ws, json.dumps({
                            'type': 'shell_output',
                            'shellId': self.shell_id,
                            'data': data,
                        }))
                except EOFError:
                    break
                except Exception as e:
                    print(f'[winpty 스트림:{self.shell_id}] {e}', flush=True)
                    break
        except asyncio.CancelledError:
            pass
        try:
            await ws.send(json.dumps({'type': 'shell_exit', 'shellId': self.shell_id, 'code': 0}))
        except Exception:
            pass

    async def _stream_pipe(self, ws):
        try:
            while self.proc and self.proc.returncode is None:
                try:
                    chunk = await asyncio.wait_for(self.proc.stdout.read(512), timeout=0.05)
                    if chunk:
                        await _safe_send(ws, json.dumps({
                            'type': 'shell_output',
                            'shellId': self.shell_id,
                            'data': chunk.decode('utf-8', errors='replace').replace('\x00', ''),
                        }))
                except asyncio.TimeoutError:
                    continue
        except asyncio.CancelledError:
            pass
        except Exception as e:
            print(f'[pipe 오류:{self.shell_id}] {e}', flush=True)
        if self.proc:
            try:
                await self.proc.wait()
            except Exception:
                pass
        try:
            await ws.send(json.dumps({
                'type': 'shell_exit',
                'shellId': self.shell_id,
                'code': self.proc.returncode if self.proc else -1,
            }))
        except Exception:
            pass

    def resize(self, cols: int, rows: int):
        if IS_WINDOWS:
            if self._win_pty:
                try:
                    if hasattr(self._win_pty, 'set_size'):
                        self._win_pty.set_size(cols, rows)
                    elif hasattr(self._win_pty, 'setwinsize'):
                        self._win_pty.setwinsize(rows, cols)
                except Exception:
                    pass
            return
        if self.master_fd is None:
            return
        try:
            import fcntl, termios, struct
            winsize = struct.pack('HHHH', rows, cols, 0, 0)
            fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, winsize)
        except Exception:
            pass

    def write_input(self, data: str):
        if IS_WINDOWS and self._win_pty:
            try:
                self._win_pty.write(data)
            except Exception as e:
                print(f'[입력오류:{self.shell_id}] {e}', flush=True)
        elif self.master_fd is not None:
            try:
                os.write(self.master_fd, data.encode('utf-8'))
            except Exception as e:
                print(f'[입력오류:{self.shell_id}] {e}', flush=True)
        elif self.proc and self.proc.stdin and self.proc.returncode is None:
            try:
                self.proc.stdin.write(data.encode('utf-8'))
            except Exception:
                pass


class Runner:
    """▶ 버튼 실행 전용 (run 메시지)."""

    def __init__(self):
        self.proc = None
        self.tmp_path = None

    async def kill(self):
        if self.proc and self.proc.returncode is None:
            try:
                self.proc.kill()
                await self.proc.wait()
            except Exception:
                pass
        if self.tmp_path:
            try:
                os.unlink(self.tmp_path)
            except Exception:
                pass
            self.tmp_path = None
        self.proc = None

    def write_input(self, data: str):
        if self.proc and self.proc.stdin and self.proc.returncode is None:
            try:
                self.proc.stdin.write(data.encode('utf-8'))
            except Exception:
                pass


# ── 파일 작업(Read/Write/Edit/Glob/Grep)용 헬퍼 ──────────────────────────
# claude-app의 Bash 경로와 동일한 root를 사용해 파일 작업을 처리한다.
# 이렇게 하면 Bash cwd와 Read/Write 대상 폴더가 항상 같은 곳을 가리킨다.
FS_IGNORE_DIRS = {'node_modules', '.git', 'dist', 'build', '__pycache__', 'venv', '.venv', '.next', 'target', '.cache', 'out'}


def _pick_folder():
    """네이티브 OS 폴더 선택 창을 열고 절대경로를 반환한다 (동기, executor에서 실행)."""
    import tkinter as tk
    from tkinter import filedialog
    root = tk.Tk()
    root.withdraw()
    root.attributes('-topmost', True)
    try:
        path = filedialog.askdirectory(title='claude-app 프로젝트 폴더 선택')
    finally:
        root.destroy()
    return path or None


def _pick_file(title: str, filetypes):
    """네이티브 OS 파일 선택 창을 열고 절대경로를 반환한다 (동기, executor에서 실행)."""
    import tkinter as tk
    from tkinter import filedialog
    root = tk.Tk()
    root.withdraw()
    root.attributes('-topmost', True)
    try:
        path = filedialog.askopenfilename(
            title=title or '파일 선택',
            filetypes=filetypes or [('모든 파일', '*.*')],
        )
    finally:
        root.destroy()
    return path or None


def _spawn_scrcpy(serial):
    """scrcpy를 companion.py와 완전히 독립된(disconnect/재시작에도 안 죽는) 프로세스로 띄운다.

    exec/shell 세션의 모든 자식 프로세스는 WS 연결 종료 시 handle()의
    finally 블록에서 kill되므로, scrcpy처럼 브라우저 탭과 무관하게 계속
    떠 있어야 하는 네이티브 GUI 창은 shells/runner에 등록하지 않고
    OS 세션에서 완전히 분리해서 스폰해야 한다.
    """
    import shutil
    import subprocess

    _refresh_env()
    scrcpy_path = shutil.which('scrcpy')
    if not scrcpy_path:
        return {'ok': False, 'error': 'scrcpy를 찾을 수 없습니다. https://github.com/Genymobile/scrcpy 에서 설치 후 PATH에 추가해주세요.'}

    cmd = [scrcpy_path] + (['-s', serial] if serial else [])
    log_path = tempfile.mktemp(suffix='.scrcpy.log')
    try:
        log_file = open(log_path, 'w', encoding='utf-8', errors='replace')
        popen_kwargs = dict(stdout=log_file, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL)
        if IS_WINDOWS:
            popen_kwargs['creationflags'] = subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP
        else:
            popen_kwargs['start_new_session'] = True
        proc = subprocess.Popen(cmd, **popen_kwargs)
    except Exception as e:
        return {'ok': False, 'error': f'scrcpy 실행 실패: {e}'}
    return {'ok': True, 'pid': proc.pid, '_proc': proc, '_log_path': log_path}


def _safe_join(root_path: str, rel_path: str) -> str:
    root_abs = os.path.abspath(root_path)
    full = os.path.abspath(os.path.join(root_abs, rel_path))
    if full != root_abs and not full.startswith(root_abs + os.sep):
        raise ValueError('경로가 작업 폴더를 벗어났습니다')
    return full


def _fs_walk(root_path: str):
    root_abs = os.path.abspath(root_path)
    results = []
    for dirpath, dirnames, filenames in os.walk(root_abs):
        dirnames[:] = [d for d in dirnames if d not in FS_IGNORE_DIRS and not d.startswith('.')]
        for fn in filenames:
            rel = os.path.relpath(os.path.join(dirpath, fn), root_abs).replace(os.sep, '/')
            results.append(rel)
    return results


async def handle(ws):
    shells = {}   # shell_id -> ShellSession
    runner = Runner()
    addr = ws.remote_address
    print(f'[연결] {addr}')
    try:
        async for raw in ws:
            try:
                msg = json.loads(raw)
            except Exception:
                continue

            t = msg.get('type')

            if t == 'run':
                await runner.kill()
                _refresh_env()

                code = msg.get('code', '')
                filename = msg.get('filename', 'script.py')
                ext = os.path.splitext(filename)[1] or '.py'

                with tempfile.NamedTemporaryFile(
                    mode='w', suffix=ext, delete=False, encoding='utf-8'
                ) as f:
                    f.write(code)
                    runner.tmp_path = f.name

                print(f'[실행] {filename}')
                await ws.send(json.dumps({'type': 'started', 'filename': filename}))

                try:
                    runner.proc = await asyncio.create_subprocess_exec(
                        sys.executable, runner.tmp_path,
                        stdin=asyncio.subprocess.PIPE,
                        stdout=asyncio.subprocess.PIPE,
                        stderr=asyncio.subprocess.STDOUT,
                        env={**os.environ},
                    )

                    while True:
                        chunk = await runner.proc.stdout.read(512)
                        if not chunk:
                            break
                        await ws.send(json.dumps({
                            'type': 'output',
                            'data': chunk.decode('utf-8', errors='replace'),
                        }))

                    await runner.proc.wait()
                    code_ret = runner.proc.returncode
                    print(f'[종료] {filename} → exit {code_ret}')
                    await ws.send(json.dumps({'type': 'exit', 'code': code_ret}))

                except Exception as e:
                    await ws.send(json.dumps({'type': 'error', 'data': str(e)}))

                finally:
                    if runner.tmp_path:
                        try:
                            os.unlink(runner.tmp_path)
                        except Exception:
                            pass
                    runner.tmp_path = None

            elif t == 'shell':
                shell_id = msg.get('shellId', 'default')
                rows = msg.get('rows', 24)
                cols = msg.get('cols', 80)
                _refresh_env()
                if shell_id in shells:
                    await shells[shell_id].kill()
                session = ShellSession(shell_id)
                shells[shell_id] = session
                print(f'[셸:{shell_id}] 시작 ({cols}x{rows})')
                await session.start(ws, rows=rows, cols=cols)

            elif t == 'shell_close':
                shell_id = msg.get('shellId')
                if shell_id and shell_id in shells:
                    await shells[shell_id].kill()
                    del shells[shell_id]
                    print(f'[셸:{shell_id}] 닫힘')

            elif t == 'resize':
                shell_id = msg.get('shellId')
                if shell_id and shell_id in shells:
                    shells[shell_id].resize(msg.get('cols', 80), msg.get('rows', 24))

            elif t == 'input':
                shell_id = msg.get('shellId')
                if shell_id and shell_id in shells:
                    shells[shell_id].write_input(msg.get('data', ''))
                elif not shell_id:
                    runner.write_input(msg.get('data', ''))
                    if runner.proc and runner.proc.stdin and runner.proc.returncode is None and IS_WINDOWS:
                        try:
                            await runner.proc.stdin.drain()
                        except Exception:
                            pass

            elif t == 'kill':
                shell_id = msg.get('shellId')
                if shell_id and shell_id in shells:
                    await shells[shell_id].kill()
                    del shells[shell_id]
                    print(f'[셸:{shell_id}] 강제종료')
                    await ws.send(json.dumps({'type': 'killed', 'shellId': shell_id}))
                else:
                    await runner.kill()
                    print('[강제 종료]')
                    await ws.send(json.dumps({'type': 'killed'}))

            elif t == 'exec':
                # 단발성 명령 실행 (에디터의 Bash 툴 / claude-app 에이전트 루프용).
                # 인터랙티브 셸(ShellSession)과 무관하게 reqId로 결과를 상관시킨다.
                req_id = msg.get('reqId')
                command = msg.get('command', '')
                cwd = msg.get('cwd') or None
                # Windows(cmd.exe)는 시스템 로컬 인코딩(CP949 등)을 사용하므로
                # UTF-8 디코딩 실패 시 시스템 기본 인코딩으로 폴백한다.
                _refresh_env()
                try:
                    proc = await asyncio.create_subprocess_shell(
                        command,
                        stdout=asyncio.subprocess.PIPE,
                        stderr=asyncio.subprocess.STDOUT,
                        cwd=cwd if cwd and os.path.isdir(cwd) else None,
                        env={**os.environ},
                    )
                    try:
                        out, _ = await asyncio.wait_for(proc.communicate(), timeout=30)
                        # Try UTF-8 first; on Windows fall back to system locale encoding
                        try:
                            text = out.decode('utf-8')
                        except UnicodeDecodeError:
                            text = out.decode(locale.getpreferredencoding(False), errors='replace')
                        if len(text) > 200_000:
                            text = text[:200_000] + '\n…(출력 잘림)'
                        await ws.send(json.dumps({
                            'type': 'exec_result', 'reqId': req_id,
                            'output': text, 'error': proc.returncode != 0,
                        }))
                    except asyncio.TimeoutError:
                        proc.kill()
                        await proc.wait()
                        await ws.send(json.dumps({
                            'type': 'exec_result', 'reqId': req_id,
                            'output': '오류: 명령 실행 시간 초과 (30s)', 'error': True,
                        }))
                except Exception as e:
                    await ws.send(json.dumps({
                        'type': 'exec_result', 'reqId': req_id,
                        'output': f'오류: {e}', 'error': True,
                    }))

            elif t == 'browse_folder':
                req_id = msg.get('reqId')
                try:
                    path = await asyncio.get_event_loop().run_in_executor(None, _pick_folder)
                    await ws.send(json.dumps({'type': 'browse_result', 'reqId': req_id, 'path': path}))
                except Exception as e:
                    await ws.send(json.dumps({
                        'type': 'browse_result', 'reqId': req_id, 'path': None,
                        'error': f'폴더 선택 창을 열 수 없습니다: {e}',
                    }))

            elif t == 'browse_file':
                req_id = msg.get('reqId')
                title = msg.get('title') or '파일 선택'
                filetypes = msg.get('filetypes')
                filetypes = [tuple(ft) for ft in filetypes] if filetypes else None
                try:
                    path = await asyncio.get_event_loop().run_in_executor(None, _pick_file, title, filetypes)
                    await ws.send(json.dumps({'type': 'browse_file_result', 'reqId': req_id, 'path': path}))
                except Exception as e:
                    await ws.send(json.dumps({
                        'type': 'browse_file_result', 'reqId': req_id, 'path': None,
                        'error': f'파일 선택 창을 열 수 없습니다: {e}',
                    }))

            elif t == 'scrcpy_start':
                req_id = msg.get('reqId')
                serial = msg.get('serial')
                try:
                    result = await asyncio.get_event_loop().run_in_executor(None, _spawn_scrcpy, serial)
                    if result.get('ok'):
                        proc = result.pop('_proc')
                        log_path = result.pop('_log_path')
                        await asyncio.sleep(1.2)
                        if proc.poll() is not None:
                            # 스폰 직후 바로 종료됨 — 로그 tail을 에러로 반환
                            tail = ''
                            try:
                                with open(log_path, 'r', encoding='utf-8', errors='replace') as f:
                                    tail = f.read()[-2000:]
                            except Exception:
                                pass
                            result = {'ok': False, 'error': f'scrcpy가 바로 종료되었습니다: {tail or "(로그 없음)"}'}
                    await ws.send(json.dumps({'type': 'scrcpy_result', 'reqId': req_id, **result}))
                except Exception as e:
                    await ws.send(json.dumps({'type': 'scrcpy_result', 'reqId': req_id, 'ok': False, 'error': str(e)}))

            elif t == 'fs_list':
                req_id = msg.get('reqId')
                root_path = msg.get('cwd') or ''
                try:
                    paths = _fs_walk(root_path) if root_path and os.path.isdir(root_path) else []
                    await ws.send(json.dumps({'type': 'fs_list_result', 'reqId': req_id, 'paths': paths}))
                except Exception as e:
                    await ws.send(json.dumps({'type': 'fs_list_result', 'reqId': req_id, 'paths': [], 'error': str(e)}))

            elif t == 'fs_read':
                req_id = msg.get('reqId')
                root_path = msg.get('cwd') or ''
                rel = msg.get('path') or ''
                try:
                    full = _safe_join(root_path, rel)
                    with open(full, 'r', encoding='utf-8', errors='replace') as f:
                        content = f.read()
                    await ws.send(json.dumps({'type': 'fs_read_result', 'reqId': req_id, 'content': content}))
                except Exception as e:
                    await ws.send(json.dumps({'type': 'fs_read_result', 'reqId': req_id, 'content': None, 'error': str(e)}))

            elif t == 'fs_write':
                req_id = msg.get('reqId')
                root_path = msg.get('cwd') or ''
                rel = msg.get('path') or ''
                content = msg.get('content', '')
                try:
                    full = _safe_join(root_path, rel)
                    os.makedirs(os.path.dirname(full), exist_ok=True)
                    with open(full, 'w', encoding='utf-8') as f:
                        f.write(content)
                    await ws.send(json.dumps({'type': 'fs_write_result', 'reqId': req_id, 'ok': True}))
                except Exception as e:
                    await ws.send(json.dumps({'type': 'fs_write_result', 'reqId': req_id, 'ok': False, 'error': str(e)}))

            elif t == 'ping':
                await ws.send(json.dumps({'type': 'pong'}))

    except Exception as e:
        print(f'[연결오류] {type(e).__name__}: {e}', flush=True)
    finally:
        for session in list(shells.values()):
            await session.kill()
        await runner.kill()
        print(f'[종료] {addr}', flush=True)


def _make_tray_icon():
    from PIL import Image, ImageDraw
    size = 64
    img = Image.new('RGBA', (size, size), (0, 0, 0, 0))
    d = ImageDraw.Draw(img)
    d.ellipse([2, 2, size - 2, size - 2], fill=(52, 120, 200))
    d.ellipse([18, 18, 46, 46], fill='white')
    d.ellipse([24, 24, 40, 40], fill=(52, 120, 200))
    return img


def _run_asyncio_in_thread(stop_event):
    async def _serve():
        async with websockets.serve(handle, HOST, PORT):
            while not stop_event.is_set():
                await asyncio.sleep(0.3)
    try:
        asyncio.run(_serve())
    except Exception as e:
        print(f'[서버오류] {e}', flush=True)


def _run_with_tray():
    import pystray
    stop_event = threading.Event()
    t = threading.Thread(target=_run_asyncio_in_thread, args=(stop_event,), daemon=True)
    t.start()

    def on_quit(icon, item):
        stop_event.set()
        icon.stop()

    icon = pystray.Icon(
        'python_companion',
        _make_tray_icon(),
        f'Python Companion  ws://{HOST}:{PORT}',
        menu=pystray.Menu(
            pystray.MenuItem('Python Companion', None, enabled=False),
            pystray.MenuItem(f'포트: {PORT}  |  연결 대기 중', None, enabled=False),
            pystray.Menu.SEPARATOR,
            pystray.MenuItem('종료', on_quit),
        ),
    )
    print('시스템 트레이에서 실행 중 (우클릭 → 종료)', flush=True)
    icon.run()
    stop_event.set()


async def main():
    async with websockets.serve(handle, HOST, PORT):
        await asyncio.Future()


def _check_dependencies():
    """선택적 의존성 상태를 확인해 시작 시 한눈에 보여준다.
    websockets는 필수라 파일 최상단에서 이미 체크됨(없으면 여기까지 못 옴)."""
    checks = [('websockets (필수)', True, '')]

    try:
        import tkinter  # noqa: F401
        checks.append(('tkinter — 폴더 선택창', True, ''))
    except ImportError:
        hint = 'python.org 설치본은 기본 포함, 커스텀/최소 설치라면 재설치 필요' if IS_WINDOWS else 'sudo apt install python3-tk (또는 배포판에 맞는 tk 패키지)'
        checks.append(('tkinter — 폴더 선택창', False, hint))

    if IS_WINDOWS:
        has_pty = False
        for mod in ('pywinpty', 'winpty'):
            try:
                __import__(mod)
                has_pty = True
                break
            except ImportError:
                continue
        checks.append(('pywinpty — 터미널 색상/크기 지원', has_pty, '' if has_pty else 'pip install pywinpty'))

    try:
        import pystray  # noqa: F401
        from PIL import Image  # noqa: F401
        checks.append(('pystray+Pillow — 시스템 트레이 실행', True, ''))
    except ImportError:
        checks.append(('pystray+Pillow — 시스템 트레이 실행', False, 'pip install pystray Pillow'))

    print('의존성 상태:', flush=True)
    for name, ok, hint in checks:
        line = f'  {"✅" if ok else "⚠️ "} {name}'
        if not ok and hint:
            line += f'  →  {hint}'
        print(line, flush=True)
    print(flush=True)


if __name__ == '__main__':
    try:
        print('━' * 40)
        print(f'Python Companion  ws://{HOST}:{PORT}')
        print(f'플랫폼: {"Windows" if IS_WINDOWS else "Unix/Linux"}')
        print('━' * 40)
        print(f'Python: {sys.executable}')
        print()
    except Exception:
        pass

    _check_dependencies()

    try:
        import pystray
        from PIL import Image
        _run_with_tray()
    except ImportError:
        print('웹 에디터에서 연결 대기 중...\nCtrl+C 로 종료\n', flush=True)
        try:
            asyncio.run(main())
        except KeyboardInterrupt:
            print('\n컴패니언 종료')