# updater.py — Windows, PyInstaller onefile 권장 (로깅/진단/언인스/백업/복원 포함)
import argparse, os, shutil, subprocess, sys, time, tempfile, uuid, datetime, csv, io, traceback
import json
import threading
import tkinter as tk
from tkinter import ttk

try:
    import winreg
except ImportError:
    winreg = None

# 기본 종료 대상 (필요시 --targets 로 덮어쓰기)
TARGETS = ["Edit_PartTimer3.exe", "main.exe"]

UNINSTALL_KEYS = [
    r"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall",
    r"SOFTWARE\WOW6432Node\Microsoft\Windows\CurrentVersion\Uninstall",
]

# ── 로깅 ─────────────────────────────────────────────────────────────────────
LOG_PATH = os.path.join(tempfile.gettempdir(), "updater_debug.log")

def log(msg: str, exc_info: bool = False):
    try:
        ts = datetime.datetime.now().isoformat(timespec="seconds")
        with open(LOG_PATH, "a", encoding="utf-8") as f:
            f.write(f"[{ts}] {msg}\n")
            if exc_info:
                f.write(traceback.format_exc())
    except Exception:
        pass

# ── GUI 진행 상황 표시 ─────────────────────────────────────────────────────
class ProgressWindow:
    """업데이트 진행 상황을 표시하는 간단한 GUI 창"""
    def __init__(self):
        self.root = None
        self.label = None
        self.progress = None
        self.detail_label = None
        self._thread = None
        
    def start(self):
        """GUI 스레드 시작"""
        self._thread = threading.Thread(target=self._run_gui, daemon=True)
        self._thread.start()
        time.sleep(0.2)  # GUI 초기화 대기
        
    def _run_gui(self):
        """GUI 메인 루프 (별도 스레드에서 실행)"""
        self.root = tk.Tk()
        self.root.title("편집알바생 업데이트 중")
        self.root.geometry("450x180")
        self.root.resizable(False, False)
        
        # 창을 항상 위에 표시
        self.root.attributes("-topmost", True)
        
        # 중앙 배치
        self.root.update_idletasks()
        width = self.root.winfo_width()
        height = self.root.winfo_height()
        x = (self.root.winfo_screenwidth() // 2) - (width // 2)
        y = (self.root.winfo_screenheight() // 2) - (height // 2)
        self.root.geometry(f'{width}x{height}+{x}+{y}')
        
        # 프레임
        frame = ttk.Frame(self.root, padding="20")
        frame.pack(fill=tk.BOTH, expand=True)
        
        # 제목
        title = ttk.Label(frame, text="업데이트를 진행하고 있습니다", 
                         font=("맑은 고딕", 12, "bold"))
        title.pack(pady=(0, 15))
        
        # 상태 메시지
        self.label = ttk.Label(frame, text="업데이트를 준비하는 중...", 
                              font=("맑은 고딕", 10))
        self.label.pack(pady=(0, 10))
        
        # 프로그레스 바
        self.progress = ttk.Progressbar(frame, mode='indeterminate', length=380)
        self.progress.pack(pady=(0, 10))
        self.progress.start(10)
        
        # 상세 정보
        self.detail_label = ttk.Label(frame, text="잠시만 기다려주세요...", 
                                     font=("맑은 고딕", 8), foreground="gray")
        self.detail_label.pack()
        
        self.root.mainloop()
        
    def update_status(self, message: str, detail: str = ""):
        """상태 메시지 업데이트"""
        if self.root and self.label:
            try:
                self.label.config(text=message)
                if detail and self.detail_label:
                    self.detail_label.config(text=detail)
                self.root.update()
            except:
                pass
                
    def close(self):
        """창 닫기"""
        if self.root:
            try:
                self.root.quit()
                self.root.destroy()
            except:
                pass

# 전역 프로그레스 윈도우
progress_window = None

# ── 유틸 ─────────────────────────────────────────────────────────────────────
def user_data_path(appdir: str) -> str:
    return os.path.join(appdir, "lib", "src", "user_data")

def tasklist_names() -> set[str]:
    try:
        out = subprocess.check_output(["tasklist", "/FO", "CSV", "/NH"], encoding="cp949", errors="replace")
        names = set()
        for line in out.splitlines():
            parts = [p.strip('"') for p in line.split(",")]
            if parts and parts[0]:
                names.add(parts[0])
        return {n.lower() for n in names}
    except Exception as e:
        log(f"tasklist read failed: {e}")
        return set()

def wait_until_killed(targets: list[str], timeout_s=10) -> bool:
    deadline = time.time() + timeout_s
    tset = {t.lower() for t in targets}
    while time.time() < deadline:
        alive = tasklist_names() & tset
        if not alive:
            return True
        time.sleep(0.4)
    log(f"still alive after kill wait: {', '.join(sorted(alive))}")
    return False

def pids_by_image(name: str) -> list[int]:
    """tasklist로 특정 이미지 이름의 PID 목록을 얻는다."""
    try:
        out = subprocess.check_output(
            ["tasklist", "/FI", f"IMAGENAME eq {name}", "/FO", "CSV", "/NH"],
            encoding="cp949", errors="replace"
        )
    except Exception:
        return []
    pids = []
    for line in out.splitlines():
        if not line.strip() or "정보:" in line:  # "정보: 일치하는 작업이 없습니다." 등 한글 메시지
            continue
        row = next(csv.reader(io.StringIO(line)))
        # row = ["Image Name","PID","Session Name","Session#","Mem Usage"]
        if len(row) >= 2 and row[0].lower() == name.lower():
            try:
                pids.append(int(row[1]))
            except ValueError:
                pass
    return pids

def kill_targets(targets: list[str]):
    log(f"kill_targets START: {targets}")
    self_pid = os.getpid()
    log(f"self_pid={self_pid}")
    # PyInstaller onefile일 때도 이 값이 "현재 실행 중인 (temp 경로의) updater.exe" PID
    for img in targets:
        log(f"Processing image: {img}")
        try:
            # 자기 자신 이미지명이라도 PID로 걸러서 본인 제외
            pids = pids_by_image(img)
            log(f"Found PIDs for {img}: {pids}")
            safe_pids = [pid for pid in pids if pid != self_pid]
            log(f"Safe PIDs for {img}: {safe_pids}")
            
            if not safe_pids:
                log(f"no killable PIDs for {img} (found={pids}, self_pid={self_pid})")
                continue
            
            # 개별 PID로 종료하면 자기 자신을 실수로 죽일 일이 없음
            cmd = ["taskkill", "/F"]  # /T 옵션 제거 (자식 프로세스까지 죽이면 위험)
            for pid in safe_pids:
                cmd.extend(["/PID", str(pid)])
            
            log(f"Executing taskkill command: {' '.join(cmd)}")
            r = subprocess.run(cmd, capture_output=True, text=True, shell=False, timeout=10)
            log(f"taskkill {img} PIDs={safe_pids} -> rc={r.returncode} out='{(r.stdout or '').strip()}' err='{(r.stderr or '').strip()}'")
        except subprocess.TimeoutExpired:
            log(f"taskkill {img} timeout")
        except Exception as e:
            log(f"taskkill {img} exception: {e}", exc_info=True)
    
    log("kill_targets: waiting for processes to die...")
    # 짧게 안정 대기
    deadline = time.time() + 10
    while time.time() < deadline:
        alive = []
        try:
            for img in targets:
                for pid in pids_by_image(img):
                    if pid != self_pid:
                        alive.append((img, pid))
            if not alive:
                break
        except Exception as e:
            log(f"Error checking alive processes: {e}")
            break
        time.sleep(0.4)
    log(f"kill wait done; alive after wait: {alive if 'alive' in locals() else []}")

def run_installer(installer_path: str, silent: bool, inno_log: str, appdir: str | None = None) -> int:
    args = [installer_path]
    if silent:
        args += ["/VERYSILENT", "/SUPPRESSMSGBOXES", "/NORESTART", "/CLOSEAPPLICATIONS", "/RESTARTAPPLICATIONS"]
    if appdir:
        args += [f'/DIR={appdir}']   # subprocess.run이 공백 자동 처리
    args += [f'/LOG={inno_log}']     # 따옴표 제거
    log(f"run_installer: args={args}")
    try:
        completed = subprocess.run(args, check=False, shell=False, close_fds=True)
        log(f"installer rc={completed.returncode}")
        return completed.returncode
    except Exception as e:
        log(f"installer exception: {e}", exc_info=True)
        return 1

def relaunch(appdir: str, appexe: str):
    import os, subprocess
    # 실행 파일 전체 경로
    full = os.path.join(appdir, appexe) if appdir else appexe
    exists = os.path.exists(full)
    log(f"relaunch: {full} exists={exists}")
    if not exists:
        log("relaunch abort: target exe not found")
        return

    # 작업폴더(cwd) 결정: appdir이 있으면 그걸 사용, 없으면 exe의 폴더
    cwd = appdir if appdir else os.path.dirname(full)
    try:
        subprocess.Popen([full], shell=False, cwd=cwd, close_fds=True)
        log("relaunch Popen OK")
    except Exception as e:
        log(f"relaunch error: {e}")

def find_uninstaller_by_appid_or_name(app_name_keywords=(), app_id=None):
    if winreg is None:
        return None
    hives = [(winreg.HKEY_LOCAL_MACHINE, "HKLM"), (winreg.HKEY_CURRENT_USER, "HKCU")]
    for hive, _ in hives:
        for base in UNINSTALL_KEYS:
            try:
                with winreg.OpenKey(hive, base) as k:
                    idx = 0
                    while True:
                        try:
                            subkey_name = winreg.EnumKey(k, idx)
                        except OSError:
                            break
                        idx += 1
                        try:
                            with winreg.OpenKey(k, subkey_name) as sk:
                                display_name = None
                                uninstall_str = None
                                try:
                                    display_name, _ = winreg.QueryValueEx(sk, "DisplayName")
                                except OSError:
                                    pass
                                try:
                                    uninstall_str, _ = winreg.QueryValueEx(sk, "UninstallString")
                                except OSError:
                                    pass
                                if app_id:
                                    try:
                                        appid_reg, _ = winreg.QueryValueEx(sk, "Inno Setup: AppId")
                                        if appid_reg and appid_reg.strip().lower() == app_id.strip().lower():
                                            if uninstall_str:
                                                return uninstall_str
                                    except OSError:
                                        pass
                                if display_name and app_name_keywords:
                                    name_l = display_name.lower()
                                    if all(kw.lower() in name_l for kw in app_name_keywords):
                                        if uninstall_str:
                                            return uninstall_str
                        except OSError:
                            continue
            except OSError:
                continue
    return None

def run_uninstaller(uninstall_cmd: str, unins_log: str) -> int:
    if not uninstall_cmd:
        return 0
    cmdline = uninstall_cmd.strip()
    extra = f' /VERYSILENT /SUPPRESSMSGBOXES /NORESTART /LOG="{unins_log}"'
    if "/LOG" not in cmdline.upper():
        cmdline = f"{cmdline} {extra}"
    log(f"run_uninstaller: {cmdline}")
    try:
        completed = subprocess.run(cmdline, check=False, shell=True, close_fds=True)
        log(f"uninstaller rc={completed.returncode}")
        return completed.returncode
    except Exception as e:
        log(f"uninstaller exception: {e}")
        return 1

def find_unins_in_appdir(appdir: str):
    if not appdir or not os.path.isdir(appdir):
        return None
    for name in os.listdir(appdir):
        if name.lower().startswith("unins") and name.lower().endswith(".exe"):
            return os.path.join(appdir, name)
    return None

def backup_user_data(appdir: str) -> str | None:
    src = user_data_path(appdir)
    if not appdir or not os.path.isdir(src):
        log("backup: no user_data to backup")
        return None
    backup_root = os.path.join(tempfile.gettempdir(), f"user_data_backup_{uuid.uuid4().hex}")
    try:
        shutil.copytree(src, backup_root, dirs_exist_ok=True)
        log(f"backup: {src} -> {backup_root}")
        return backup_root
    except Exception as e:
        log(f"backup error: {e}")
        return None

def restore_user_data(appdir: str, backup_dir: str | None):
    if not backup_dir or not os.path.isdir(backup_dir) or not appdir:
        log("restore: skip (no backup)")
        return
    dst = user_data_path(appdir)
    os.makedirs(dst, exist_ok=True)
    copied = 0
    for root, dirs, files in os.walk(backup_dir):
        rel = os.path.relpath(root, backup_dir)
        target = os.path.join(dst, rel) if rel != "." else dst
        os.makedirs(target, exist_ok=True)
        for d in dirs:
            os.makedirs(os.path.join(target, d), exist_ok=True)
        for f in files:
            src_file = os.path.join(root, f)
            dst_file = os.path.join(target, f)
            try:
                shutil.copy2(src_file, dst_file); copied += 1
            except Exception as e:
                log(f"restore copy error {src_file} -> {dst_file}: {e}")
    log(f"restore: copied {copied} files to {dst}")
    try:
        shutil.rmtree(backup_dir, ignore_errors=True)
    except Exception:
        pass

def find_uninstaller_by_appid_fuzzy(appid: str) -> str | None:
    """
    Inno 언인스키는 보통 '<AppId>_is1'. AppId가 GUID든 문자열이든 모두 시도.
    """
    if not appid or winreg is None:
        return None
    # 시도할 서픽스 후보 만들기
    cand_suffixes = []
    a = appid.strip()
    cand_suffixes.append((a + "_is1").lower())
    if a.startswith("{") and a.endswith("}"):
        naked = a[1:-1]
        cand_suffixes.append((naked + "_is1").lower())
    else:
        cand_suffixes.append(("{"+a+"}" + "_is1").lower())

    # 레지스트리 열거
    hives = [(winreg.HKEY_LOCAL_MACHINE, "HKLM"), (winreg.HKEY_CURRENT_USER, "HKCU")]
    for hive, _ in hives:
        for base in UNINSTALL_KEYS:
            try:
                with winreg.OpenKey(hive, base) as k:
                    i = 0
                    while True:
                        try:
                            sub = winreg.EnumKey(k, i); i += 1
                        except OSError:
                            break
                        sub_l = sub.lower()
                        if any(sub_l.endswith(suf) for suf in cand_suffixes):
                            try:
                                with winreg.OpenKey(k, sub) as sk:
                                    uninstall_str, _ = winreg.QueryValueEx(sk, "UninstallString")
                                    if uninstall_str:
                                        return uninstall_str
                            except OSError:
                                pass
            except OSError:
                continue
    return None

def read_manifest_near(installer_path: str) -> dict:
    """
    설치본(Setup_*.exe) 옆에 '{installer}.manifest.json'이 있으면 읽는다.
    거기에 app_id, app_name, default_dirname 등이 들어있다고 가정.
    """
    try:
        base = os.path.splitext(installer_path)[0]  # .exe 제거
        cand = base + ".manifest.json"
        if os.path.exists(cand):
            with open(cand, "r", encoding="utf-8") as f:
                return json.load(f)
    except Exception as e:
        log(f"manifest read error: {e}")
    return {}

def resolve_appdir_from_registry(appid: str, app_name_keywords: tuple[str, ...]) -> str:
    """
    InstallLocation 또는 언인스 exe 경로의 상위 폴더로 appdir 추정.
    """
    if winreg is None:
        return ""
    def _search_entry():
        hives = [(winreg.HKEY_LOCAL_MACHINE, "HKLM"), (winreg.HKEY_CURRENT_USER, "HKCU")]
        for hive, _ in hives:
            for base in UNINSTALL_KEYS:
                try:
                    with winreg.OpenKey(hive, base) as k:
                        i = 0
                        while True:
                            try:
                                sub = winreg.EnumKey(k, i); i += 1
                            except OSError:
                                break
                            try:
                                with winreg.OpenKey(k, sub) as sk:
                                    # AppId 우선
                                    if appid:
                                        try:
                                            v, _ = winreg.QueryValueEx(sk, "Inno Setup: AppId")
                                            if v and v.strip().lower() == appid.strip().lower():
                                                return sk
                                        except OSError:
                                            pass
                                    # 키워드로 DisplayName 매칭
                                    if app_name_keywords:
                                        try:
                                            dn, _ = winreg.QueryValueEx(sk, "DisplayName")
                                            if dn and all(kw.lower() in dn.lower() for kw in app_name_keywords):
                                                return sk
                                        except OSError:
                                            pass
                            except OSError:
                                pass
                except OSError:
                    continue
        return None

    sk = _search_entry()
    if not sk:
        return ""
    # InstallLocation → 최선
    try:
        il, _ = winreg.QueryValueEx(sk, "InstallLocation")
        if il and os.path.isdir(il):
            return il
    except OSError:
        pass
    # UninstallString 경로에서 상위 폴더 추정
    try:
        us, _ = winreg.QueryValueEx(sk, "UninstallString")
        if us:
            # 예: "C:\Program Files\...\unins000.exe" /SILENT ...
            # 첫 따옴표 블록 추출
            if us.startswith('"'):
                p = us.split('"')
                if len(p) >= 3:
                    unins = p[1]
                else:
                    unins = us.strip().split()[0]
            else:
                unins = us.strip().split()[0]
            if os.path.isfile(unins):
                return os.path.dirname(unins)
    except OSError:
        pass
    return ""

# ── 메인 ─────────────────────────────────────────────────────────────────────
def main():

    try:
        os.chdir(tempfile.gettempdir())   # _MEI가 아닌 안전한 폴더로 CWD 이동
    except Exception:
        pass

    global LOG_PATH  # ← FIX: 전역 재할당 전에 선언

    ap = argparse.ArgumentParser(description="Minimal updater (Windows, PyInstaller onefile)")
    ap.add_argument("--installer", required=True)
    ap.add_argument("--appexe",   required=True)
    ap.add_argument("--appdir",   default="")
    ap.add_argument("--nosilent", action="store_true")
    ap.add_argument("--appid",    default="")
    ap.add_argument("--appnamekw", nargs="*", default=["Edit", "PartTimer"])
    ap.add_argument("--targets",  default="", help="쉼표구분 프로세스명 목록으로 TARGETS 대체 (예: a.exe,b.exe)")
    ap.add_argument("--log",      default=LOG_PATH, help="updater log path")
    a = ap.parse_args()

    # 로그 경로 설정을 먼저 적용
    LOG_PATH = a.log

    log("===== updater start =====")
    log(f"args: installer='{a.installer}', appexe='{a.appexe}', appdir='{a.appdir}', nosilent={a.nosilent}, appid='{a.appid}', appnamekw={a.appnamekw}, targets='{a.targets}'")
    
    # installer 변수 먼저 정의
    installer = os.path.abspath(a.installer)
    appdir    = os.path.abspath(a.appdir) if a.appdir else ""

    # ---- 자동결정 (appid/appdir) ----
    mf = read_manifest_near(installer)
    # 우선순위: CLI > 매니페스트 > 빈값
    appid = a.appid.strip() or str(mf.get("app_id", "")).strip()
    if not appdir:
        # 매니페스트에 default_dirname 있으면 Program Files 하위 기본값 시도
        default_dir = str(mf.get("default_dirname", "")).strip()
        if default_dir:
            pf = os.environ.get("ProgramFiles", r"C:\Program Files")
            cand = os.path.join(pf, default_dir)
            if os.path.isdir(cand):
                appdir = cand
        # 그래도 못 찾으면 레지스트리에서 설치경로 역추적
        if not appdir:
            appdir = resolve_appdir_from_registry(appid, tuple(a.appnamekw))
    # 최종값 로깅
    log(f"auto-decided: appid='{appid}', appdir='{appdir}'")
    inno_log  = os.path.join(tempfile.gettempdir(), "inno_install.log")
    unins_log = os.path.join(tempfile.gettempdir(), "inno_uninstall.log")
    log(f"resolved: installer='{installer}', appdir='{appdir}', inno_log='{inno_log}', unins_log='{unins_log}'")

    # 런타임 대상 프로세스 조정
    targets = TARGETS.copy()
    if a.targets.strip():
        targets = [t.strip() for t in a.targets.split(",") if t.strip()]
        log(f"targets override -> {targets}")

    if not os.path.exists(installer):
        log("installer not found -> abort(2)")
        return 2

    # GUI 시작 (항상 표시)
    global progress_window
    progress_window = ProgressWindow()
    progress_window.start()
    log("Progress window started")

    try:
        # 1) 프로세스 종료
        if progress_window:
            progress_window.update_status("프로그램을 종료하는 중...", 
                                         "실행 중인 프로세스를 안전하게 종료합니다")
        kill_targets(targets)

        # 2) user_data 백업
        if progress_window:
            progress_window.update_status("데이터를 백업하는 중...", 
                                         "사용자 데이터를 안전하게 백업합니다")
        backup_dir = backup_user_data(appdir)

        # 3) 기존 버전 언인스톨
        if progress_window:
            progress_window.update_status("이전 버전을 제거하는 중...", 
                                         "기존 설치 파일을 제거합니다")
        uninstall_cmd = None
        
        # appid로 먼저 시도
        if appid:
            uninstall_cmd = find_uninstaller_by_appid_fuzzy(appid)
            if uninstall_cmd:
                log(f"uninstall_cmd (appid fuzzy) = {uninstall_cmd}")

        # 키워드 검색으로 시도
        if not uninstall_cmd and (a.appid or (a.appnamekw and len(a.appnamekw) > 0)):
            uninstall_cmd = find_uninstaller_by_appid_or_name(tuple(a.appnamekw), a.appid)
            if uninstall_cmd:
                log(f"uninstall_cmd (registry) = {uninstall_cmd}")
        if not uninstall_cmd and appdir:
            unins_path = find_unins_in_appdir(appdir)
            if unins_path:
                uninstall_cmd = f'"{unins_path}"'
                log(f"uninstall_cmd (appdir) = {uninstall_cmd}")

        unins_rc = 0
        if uninstall_cmd:
            unins_rc = run_uninstaller(uninstall_cmd, unins_log)
            time.sleep(0.7)

        # 4) 새 버전 설치
        if progress_window:
            progress_window.update_status("새 버전을 설치하는 중...", 
                                         "업데이트 파일을 설치합니다 (약 30초 소요)")
        rc = run_installer(installer, silent=not a.nosilent, inno_log=inno_log, appdir=appdir)

        # 5) user_data 복원
        if progress_window:
            progress_window.update_status("데이터를 복원하는 중...", 
                                         "백업한 사용자 데이터를 복원합니다")
        restore_user_data(appdir, backup_dir)
        time.sleep(0.15)  # 150ms 정도

        # 6) 메인 앱 재실행
        if progress_window:
            progress_window.update_status("업데이트 완료!", 
                                         "프로그램을 다시 시작합니다...")
        time.sleep(0.5)  # 메시지 표시 시간
        
        relaunch(appdir, a.appexe)

        log(f"return code: install={rc}, uninstall={unins_rc}")
        log("===== updater end =====")
        return rc if rc != 0 else unins_rc or 0
    
    finally:
        # GUI 종료
        if progress_window:
            time.sleep(0.3)  # 완료 메시지 표시
            progress_window.close()
            log("Progress window closed")

if __name__ == "__main__":
    sys.exit(main())