"""
mokai Worker v2.4  —  Python-Free Executable Edition
PyInstaller --onefile --noconsole 로 빌드된 단일 실행파일.
Python 미설치 환경에서도 워커 본체가 동작.
run_python/install_packages/assign_role 같은 'Python 필요' 명령은
시스템 Python을 자동 탐색하여 사용하고, 없으면 명확한 에러 반환.
"""
import os, sys, json, hashlib, platform, socket, threading, time, subprocess, glob
import base64, logging, shutil

FROZEN  = getattr(sys, "frozen", False)
EXE_PATH = sys.executable if FROZEN else os.path.abspath(__file__)
BASE_DIR = os.path.dirname(EXE_PATH)
os.makedirs(BASE_DIR,                        exist_ok=True)
os.makedirs(os.path.join(BASE_DIR, "tasks"), exist_ok=True)

_handlers = [logging.FileHandler(os.path.join(BASE_DIR, "worker.log"), encoding="utf-8")]
if sys.stdout:
    try:
        _handlers.insert(0, logging.StreamHandler(sys.stdout))
    except Exception:
        pass
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [mokai] %(message)s",
    datefmt="%H:%M:%S",
    handlers=_handlers,
)
log = logging.getLogger()

import requests
from flask import Flask, request as freq, jsonify

CONFIG_PATH = os.path.join(BASE_DIR, "config.json")

def load_config():
    if os.path.exists(CONFIG_PATH):
        try:
            with open(CONFIG_PATH, encoding="utf-8-sig") as f:
                return json.load(f)
        except Exception as e:
            log.error(f"config.json 파싱 오류: {e}")
    return {}

cfg = load_config()

MY_PORT   = int(cfg.get("my_port",   os.environ.get("MY_PORT", "8000")))
WORKER_ID = cfg.get("worker_id", socket.gethostname())
API_KEY   = cfg.get("api_key", "")
VERSION   = "2.5"

DISCOVERY_PORT  = 9401
DISCOVERY_MAGIC = b"MOKAI_REG_DISCOVER?"

def discover_registry(timeout=2.0, tries=2):
    """LAN UDP 브로드캐스트로 레지스트리 서버 자동 발견 → 'ip:port' 또는 None.
       도메인/고정IP 없이 같은 LAN 안의 레지스트리를 환경 무관하게 찾는다."""
    for _ in range(tries):
        s = None
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
            s.settimeout(timeout)
            s.sendto(DISCOVERY_MAGIC, ("255.255.255.255", DISCOVERY_PORT))
            data, _ = s.recvfrom(256)
            if data.startswith(b"MOKAI_REG="):
                return data.decode("utf-8", "ignore").split("=", 1)[1].strip()
        except Exception:
            pass
        finally:
            if s is not None:
                s.close()
    return None

REGISTRY_URLS = [
    "https://work.mok.ai.kr/api/workers",
    "http://172.30.1.23:9400/api/workers",
]
REGISTRY_REGISTER_URLS = [
    "https://work.mok.ai.kr/api/workers/register",
    "http://172.30.1.23:9400/api/workers/register",
]

# 부팅 시 LAN 브로드캐스트로 레지스트리 자동 발견 → 발견되면 LAN을 1순위, 도메인은 폴백
_disc = discover_registry()
if _disc:
    REGISTRY_URLS.insert(0,          f"http://{_disc}/api/workers")
    REGISTRY_REGISTER_URLS.insert(0, f"http://{_disc}/api/workers/register")
    log.info(f"[Discovery] 레지스트리 자동 발견 → {_disc} (1순위 등록, 도메인은 폴백)")
else:
    log.info("[Discovery] LAN 응답 없음 — 도메인(work.mok.ai.kr) 사용")

_kim_ip_cache = {"ip": cfg.get("kim_ip", ""), "ts": 0.0}
_kim_ip_lock  = threading.Lock()

_sys_python_cache = {"path": None, "checked": False}

def find_system_python():
    """시스템 Python 탐색. exe 빌드 환경에서 sys.executable이 자기 자신이므로 별도 탐색."""
    if _sys_python_cache["checked"]:
        return _sys_python_cache["path"]
    candidates = []
    p = shutil.which("python")
    if p: candidates.append(p)
    p = shutil.which("python3")
    if p: candidates.append(p)
    p = shutil.which("py")
    if p: candidates.append(p + " -3")
    for base in [os.environ.get("LOCALAPPDATA", ""), "C:\\Python311", "C:\\Python312", "C:\\Python310"]:
        if not base: continue
        for sub in ["Programs\\Python\\Python311", "Programs\\Python\\Python312",
                    "Programs\\Python\\Python310", ""]:
            cand = os.path.join(base, sub, "python.exe")
            if os.path.exists(cand): candidates.append(cand)
    for cand in candidates:
        try:
            exe = cand.split(" ")[0]
            args = cand.split(" ")[1:] + ["--version"]
            r = subprocess.run([exe] + args, capture_output=True, text=True,
                               timeout=5, creationflags=0x08000000)
            if r.returncode == 0 and "Python" in (r.stdout + r.stderr):
                _sys_python_cache["path"] = cand
                _sys_python_cache["checked"] = True
                log.info(f"[python] 시스템 Python 탐지: {cand}")
                return cand
        except Exception:
            continue
    _sys_python_cache["checked"] = True
    log.info("[python] 시스템 Python 없음 — run_python/install_packages/assign_role 비활성화")
    return None

def py_cmd(extra_args):
    """[exe, args...] 형태로 반환. system python이 없으면 None."""
    p = find_system_python()
    if not p: return None
    parts = p.split(" ")
    return parts + extra_args

def get_kim_ip(force=False):
    with _kim_ip_lock:
        now = time.time()
        if not force and _kim_ip_cache["ip"] and now - _kim_ip_cache["ts"] < 1800:
            return _kim_ip_cache["ip"]
        for url in REGISTRY_URLS:
            try:
                r = requests.get(url, timeout=5)
                if r.status_code == 200:
                    workers = r.json().get("workers", {})
                    ip = workers.get("kim", {}).get("ip", "")
                    if ip:
                        _kim_ip_cache["ip"] = ip
                        _kim_ip_cache["ts"] = now
                        return ip
            except Exception as e:
                log.debug(f"[Registry] {url} 조회 실패: {e}")
        if _kim_ip_cache["ip"]:
            return _kim_ip_cache["ip"]
        return ""

app = Flask(__name__)

# Flask 개발서버 빨간 경고 배너 + werkzeug 요청 로그 숨김 (콘솔 깔끔하게)
try:
    import flask.cli
    flask.cli.show_server_banner = lambda *a, **k: None
except Exception:
    pass
logging.getLogger("werkzeug").setLevel(logging.ERROR)

def get_specs():
    cpu_pct = ram_total_gb = ram_free_gb = disk_free_gb = 0
    try:
        import psutil
        cpu_pct      = psutil.cpu_percent(interval=0.5)
        ram          = psutil.virtual_memory()
        disk         = psutil.disk_usage("C:/")
        ram_total_gb = round(ram.total     / 1e9, 1)
        ram_free_gb  = round(ram.available / 1e9, 1)
        disk_free_gb = round(disk.free     / 1e9, 1)
    except Exception:
        pass

    gpu = "없음"
    try:
        gpu = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=name,memory.total",
             "--format=csv,noheader,nounits"],
            text=True, timeout=3,
            stderr=subprocess.DEVNULL,
            creationflags=0x08000000).strip()
    except Exception:
        pass

    ip = "127.0.0.1"
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
    except Exception:
        pass

    return {
        "worker_id":    WORKER_ID,
        "ip":           ip,
        "port":         MY_PORT,
        "os":           platform.system() + " " + platform.release(),
        "cpu":          platform.processor()[:60],
        "cpu_cores":    os.cpu_count() or 1,
        "cpu_pct":      cpu_pct,
        "ram_total_gb": ram_total_gb,
        "ram_free_gb":  ram_free_gb,
        "disk_free_gb": disk_free_gb,
        "gpu":          gpu,
        "python":       (find_system_python() or "없음"),
        "version":      VERSION,
        "frozen":       FROZEN,
    }

current = {"role": None, "proc": None, "started": None}

def get_status():
    busy = (current["proc"] is not None and current["proc"].poll() is None)
    return {**get_specs(), "role": current["role"], "busy": busy, "alive": True}

def _stop_internal():
    proc = current["proc"]
    if proc is not None and proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
    current["role"] = current["proc"] = current["started"] = None

@app.route("/status", methods=["GET"])
def route_status():
    return jsonify(get_status())

@app.route("/health", methods=["GET"])
def route_health():
    return jsonify({"ok": True, "version": VERSION, "worker_id": WORKER_ID})

def _auth_ok():
    if not API_KEY:
        return True
    if freq.headers.get("X-Forwarded-For") or freq.headers.get("Cf-Connecting-Ip"):
        return freq.headers.get("X-API-Key") == API_KEY
    remote_ip = (freq.remote_addr or "").strip()
    if (remote_ip.startswith("172.30.") or remote_ip.startswith("192.168.")
            or remote_ip.startswith("10.") or remote_ip in ("127.0.0.1", "::1")):
        return True
    return freq.headers.get("X-API-Key") == API_KEY

@app.route("/command", methods=["POST"])
def route_command():
    if not _auth_ok():
        return jsonify({"error": "unauthorized"}), 401
    raw = freq.get_data()
    try:
        data = json.loads(raw.decode("utf-8")) if raw else {}
    except (UnicodeDecodeError, json.JSONDecodeError):
        try:
            data = json.loads(raw.decode("cp949")) if raw else {}
        except Exception:
            data = {}
    t = data.get("type", "")
    if t == "assign_role":      return cmd_assign_role(data)
    if t == "run_shell":        return cmd_run_shell(data)
    if t == "run_python":       return cmd_run_python(data)
    if t == "stop":             return cmd_stop()
    if t == "read_file":        return cmd_read_file(data)
    if t == "write_file":       return cmd_write_file(data)
    if t == "list_files":       return cmd_list_files(data)
    if t == "screenshot":       return cmd_screenshot()
    if t == "install_packages": return cmd_install_packages(data)
    if t == "ping":             return jsonify({"pong": True, "worker_id": WORKER_ID, "version": VERSION})
    return jsonify({"error": f"알 수 없는 명령: {t}"}), 400

def cmd_assign_role(data):
    role, script, pkgs = data.get("role", "WORKER"), data.get("script", ""), data.get("packages", [])
    _stop_internal()
    pip_cmd = py_cmd(["-m", "pip", "install", "--quiet"] + pkgs) if pkgs else None
    run_cmd = py_cmd(None)
    if (pkgs and not pip_cmd) or (script and not run_cmd):
        return jsonify({"ok": False, "error": "시스템 Python 미설치 — run_python/assign_role 불가"})
    if pip_cmd:
        subprocess.run(pip_cmd, capture_output=True, creationflags=0x08000000)
    if script:
        path = os.path.join(BASE_DIR, "tasks", f"{role.lower()}_task.py")
        with open(path, "w", encoding="utf-8") as f:
            f.write(script)
        proc_cmd = py_cmd([path])
        proc = subprocess.Popen(proc_cmd, creationflags=0x08000000)
        current["role"] = role
        current["proc"] = proc
        current["started"] = time.time()
        log.info(f"역할 배정: {role}")
    return jsonify({"ok": True, "role": role})

def _decode_bytes(b):
    if not b:
        return ""
    for enc in ("utf-8", "cp949", "latin-1"):
        try:
            return b.decode(enc)
        except UnicodeDecodeError:
            continue
    return b.decode("latin-1")

def cmd_run_shell(data):
    cmd, timeout = data.get("cmd", ""), int(data.get("timeout", 30))
    try:
        result = subprocess.run(cmd, shell=True, capture_output=True,
                                timeout=timeout, creationflags=0x08000000)
        return jsonify({"ok": True, "stdout": _decode_bytes(result.stdout),
                        "stderr": _decode_bytes(result.stderr), "rc": result.returncode})
    except subprocess.TimeoutExpired:
        return jsonify({"ok": False, "error": "timeout"})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)})

def cmd_run_python(data):
    code = data.get("code", "")
    path = os.path.join(BASE_DIR, "tmp_exec.py")
    pcmd = py_cmd(["-X", "utf8", path])
    if not pcmd:
        return jsonify({"ok": False, "error": "시스템 Python 미설치 — run_python 불가. install_python:true 옵션 사용 권장"})
    try:
        with open(path, "w", encoding="utf-8") as f:
            f.write(code)
        result = subprocess.run(pcmd, capture_output=True, text=True, timeout=60,
                                encoding="utf-8", errors="ignore",
                                creationflags=0x08000000)
        return jsonify({"ok": True, "stdout": result.stdout, "stderr": result.stderr})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)})

def cmd_stop():
    _stop_internal()
    return jsonify({"ok": True, "stopped": True})

def cmd_read_file(data):
    path = data.get("path", "")
    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            return jsonify({"ok": True, "content": f.read()})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)})

def cmd_write_file(data):
    path, content = data.get("path", ""), data.get("content", "")
    try:
        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        with open(path, "r", encoding="utf-8") as f:
            written = f.read()
        ok = (written == content)
        sha8 = hashlib.sha256(content.encode("utf-8")).hexdigest()[:8]
        return jsonify({"ok": ok, "size": len(content), "sha256": sha8,
                        "error": "내용 불일치" if not ok else ""})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)})

def cmd_list_files(data):
    directory, pattern = data.get("directory", BASE_DIR), data.get("pattern", "*")
    try:
        files = glob.glob(os.path.join(directory, "**", pattern), recursive=True)[:200]
        return jsonify({"ok": True, "files": files})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)})

def cmd_screenshot():
    try:
        from PIL import ImageGrab
        img  = ImageGrab.grab()
        path = os.path.join(BASE_DIR, "screen.png")
        img.save(path)
        with open(path, "rb") as f:
            b64 = base64.b64encode(f.read()).decode()
        return jsonify({"ok": True, "image_b64": b64})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)})

def cmd_install_packages(data):
    pkgs = data.get("packages", [])
    if not pkgs:
        return jsonify({"ok": False, "error": "packages 목록 없음"})
    pcmd = py_cmd(["-m", "pip", "install"] + pkgs)
    if not pcmd:
        return jsonify({"ok": False, "error": "시스템 Python 미설치 — install_packages 불가"})
    try:
        result = subprocess.run(pcmd, capture_output=True, text=True,
                                encoding="utf-8", errors="ignore",
                                creationflags=0x08000000)
        return jsonify({"ok": True, "stdout": result.stdout, "stderr": result.stderr})
    except Exception as e:
        return jsonify({"ok": False, "error": str(e)})

def register_to_hub():
    specs     = get_specs()
    services  = cfg.get("services", [f"worker_api:{MY_PORT}"])
    subdomain = cfg.get("tunnel_subdomain", "")
    payload   = {
        "name":     cfg.get("worker_name", WORKER_ID),
        "ip":       specs["ip"],
        "services": services,
    }
    if subdomain:
        payload["public_url"] = f"https://{subdomain}"
    for url in REGISTRY_REGISTER_URLS:
        try:
            r = requests.post(url, json=payload, timeout=8)
            if r.status_code == 200:
                log.info(f"[Registry] 등록 완료 → {specs['ip']}")
                return True
        except Exception as e:
            log.warning(f"[Registry] {url} 실패: {e}")
    log.error("[Registry] 등록 전체 실패 — 네트워크 확인 필요")
    return False

def register_startup_task():
    """부팅(로그온) 시 워커 자동시작을 작업스케줄러에 등록.
       - 단일 exe(frozen)   : exe 를 직접 실행
       - 동봉형/소스(비frozen): pythonw.exe -X utf8 worker.py 로 실행
         (구버전은 .py 를 Command 에 직접 넣어 부팅 시 실행이 안 됐음
          → 재부팅 후 워커가 안 뜨던 원인. 본 분기로 자동 교정한다.)
    """
    workdir   = BASE_DIR
    task_name = f"mokai_worker_{WORKER_ID}"

    if FROZEN:
        cmd_exe  = EXE_PATH
        cmd_args = ""
    else:
        pyw = sys.executable
        if pyw.lower().endswith("python.exe"):
            cand = pyw[:-len("python.exe")] + "pythonw.exe"
            if os.path.exists(cand):
                pyw = cand
        cmd_exe  = pyw
        cmd_args = f'-X utf8 "{EXE_PATH}"'   # EXE_PATH = worker.py 절대경로

    # 기존 등록이 '현재 올바른 실행 대상'을 가리키면 그대로 두고, 아니면 덮어쓴다.
    # (구버전이 .py 를 Command 로 넣어둔 잘못된 항목을 자동 교정)
    try:
        q = subprocess.run(["schtasks", "/query", "/tn", task_name, "/v", "/fo", "LIST"],
                           capture_output=True, creationflags=0x08000000)
        if q.returncode == 0:
            existing = _decode_bytes(q.stdout)
            if FROZEN:
                ok = EXE_PATH in existing
            else:
                ok = ("pythonw" in existing.lower()) and (EXE_PATH in existing)
            if ok:
                log.info(f"[mokai] 자동시작 이미 올바르게 등록됨 → {task_name}")
                return
            log.info(f"[mokai] 자동시작 항목 교정(재등록) → {task_name}")
    except Exception:
        pass

    args_xml  = f"\n      <Arguments>{cmd_args}</Arguments>" if cmd_args else ""
    task_xml  = f"""<?xml version="1.0" encoding="UTF-16"?>
<Task version="1.2" xmlns="http://schemas.microsoft.com/windows/2004/02/mit/task">
  <RegistrationInfo><Description>mokai Worker v{VERSION} 자동시작</Description></RegistrationInfo>
  <Triggers><LogonTrigger><Enabled>true</Enabled></LogonTrigger></Triggers>
  <Settings>
    <MultipleInstancesPolicy>IgnoreNew</MultipleInstancesPolicy>
    <DisallowStartIfOnBatteries>false</DisallowStartIfOnBatteries>
    <StopIfGoingOnBatteries>false</StopIfGoingOnBatteries>
    <ExecutionTimeLimit>PT0S</ExecutionTimeLimit>
    <Hidden>true</Hidden>
  </Settings>
  <Actions>
    <Exec>
      <Command>{cmd_exe}</Command>{args_xml}
      <WorkingDirectory>{workdir}</WorkingDirectory>
    </Exec>
  </Actions>
</Task>"""
    xml_path = os.path.join(workdir, "_worker_task.xml")
    try:
        with open(xml_path, "w", encoding="utf-16") as f:
            f.write(task_xml)
        r = subprocess.run(["schtasks", "/create", "/tn", task_name, "/xml", xml_path, "/f"],
                           capture_output=True, creationflags=0x08000000)
        if r.returncode == 0:
            log.info(f"[mokai] 자동시작 등록 완료 → {task_name}")
        else:
            log.warning(f"[Startup] 등록 실패: {_decode_bytes(r.stderr)[:80]}")
    except Exception as e:
        log.warning(f"[Startup] 오류: {e}")
    finally:
        try:
            os.remove(xml_path)
        except Exception:
            pass

def hub_refresh_loop():
    time.sleep(120)
    while True:
        try:
            register_to_hub()
            get_kim_ip(force=True)
        except Exception as e:
            log.warning(f"[Registry] 갱신 오류: {e}")
        time.sleep(1800)

def ensure_firewall():
    if platform.system() != "Windows":
        return
    cmds = [
        f"Remove-NetFirewallRule -Name 'mokai-worker' -ErrorAction SilentlyContinue; "
        f"New-NetFirewallRule -Name 'mokai-worker' -DisplayName 'mokai Worker API' "
        f"-Enabled True -Direction Inbound -Protocol TCP -Action Allow -LocalPort {MY_PORT} -Profile Any",
        "Remove-NetFirewallRule -Name 'mokai-icmp' -ErrorAction SilentlyContinue; "
        "New-NetFirewallRule -Name 'mokai-icmp' -DisplayName 'mokai ICMP Allow' "
        "-Enabled True -Direction Inbound -Protocol ICMPv4 -Action Allow -Profile Any",
    ]
    for cmd in cmds:
        try:
            subprocess.run(
                ["powershell", "-NoProfile", "-NonInteractive", "-Command", cmd],
                capture_output=True, timeout=10, creationflags=0x08000000,
            )
        except Exception:
            pass
    log.info("[Firewall] 방화벽 규칙 적용 완료 (TCP %d + ICMP, all profiles)", MY_PORT)



def cmd_poll_loop():
    """아웃바운드 명령 폴링 — 큐에서 받은 명령을 로컬 /command로 재실행 후 결과 업로드."""
    base = cfg.get("cmd_queue_url", "https://work.mok.ai.kr").rstrip("/")
    token = cfg.get("cmd_poll_token", API_KEY)
    interval = int(cfg.get("cmd_poll_interval", 5))
    exec_to = int(cfg.get("cmd_exec_timeout", 120))
    local = f"http://127.0.0.1:{MY_PORT}/command"
    log.info(f"[cmdq] poll loop start: {base} as {WORKER_ID} every {interval}s")
    while True:
        try:
            r = requests.get(f"{base}/api/cmd/poll", params={"worker": WORKER_ID},
                             headers={"X-API-Key": token}, timeout=15)
            c = r.json().get("cmd") if r.ok else None
            if c and isinstance(c.get("cmd"), dict):
                try:
                    lr = requests.post(local, json=c["cmd"],
                                       headers={"X-API-Key": API_KEY}, timeout=exec_to)
                    try: result = lr.json()
                    except Exception: result = {"status": lr.status_code, "text": lr.text[:2000]}
                except Exception as ex:
                    result = {"error": f"local exec failed: {ex}"}
                try:
                    requests.post(f"{base}/api/cmd/result",
                                  json={"worker": WORKER_ID, "id": c["id"], "result": result},
                                  headers={"X-API-Key": token}, timeout=30)
                except Exception as ex:
                    log.warning(f"[cmdq] result upload failed: {ex}")
                continue
        except Exception:
            pass
        time.sleep(interval)


def main():
    # 단일 실행 보장 — 같은 포트에서 워커가 이미 돌고 있으면 즉시 종료 (중복 방지)
    _probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    _probe.settimeout(1)
    if _probe.connect_ex(("127.0.0.1", MY_PORT)) == 0:
        _probe.close()
        log.info(f"[mokai] 포트 {MY_PORT} 이미 사용 중 — 중복 실행 방지로 종료합니다")
        sys.exit(0)
    _probe.close()

    worker_name = cfg.get("worker_name", WORKER_ID)
    subdomain   = cfg.get("tunnel_subdomain", "")
    log.info(f"mokai Worker v{VERSION} — {worker_name} (포트 {MY_PORT}) [frozen={FROZEN}]")
    log.info(f"  실행파일: {EXE_PATH}")
    log.info(f"  레지스트리: {REGISTRY_URLS[0]}")
    if subdomain:
        log.info(f"  공개 URL:   https://{subdomain}")
    ensure_firewall()
    register_startup_task()
    register_to_hub()
    threading.Thread(target=hub_refresh_loop, daemon=True).start()
    if cfg.get("cmd_poll_enabled"):
        threading.Thread(target=cmd_poll_loop, daemon=True).start()
        log.info("[cmdq] polling enabled")
    app.run(host="0.0.0.0", port=MY_PORT, debug=False, use_reloader=False)


if __name__ == "__main__":
    main()