#!/usr/bin/env python3
"""
Upload an Agent YAML to AgentBackend and execute it via SSE.

Flow (SKILL.md §0.6 — 경로 A):
  0. GET  {api_base}/workspaces/lookup?user_id=...    (workspace_name → UUID 해석)
  1. POST {api_base}/upload-agent?user_id=...        (multipart YAML 등록)
  2. POST {api_base}/sse/agent/{name}/start           (SSE 실행)
  3. stage_complete 이벤트 수신 시 자동 confirm       (CI 무인 실행, 재시도 포함)
  4. 연결 끊김 시 /sse/agent/{name}/reconnect/{sid}   (재연결)
  5. stage_error / --max-runtime 초과 / SIGTERM 시 세션 cancel 후 종료

Exit code: 0 = execution_complete / 1 = 그 외 (stage_error, stopped, timeout, ...)

--output-dir 에 저장되는 파일:
  events.jsonl      — 수신한 모든 SSE 이벤트 (한 줄당 1개)
  final_output.txt  — execution_complete 이벤트의 final_outputs (stage 별 출력 dict)
  summary.md        — 실행 요약 (GITHUB_STEP_SUMMARY 에도 기록)
"""

from __future__ import annotations

import argparse
import json
import os
import signal
import sys
import time
import unicodedata
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import quote

import httpx
from httpx_sse import SSEError, connect_sse

PRINT_TRUNCATE = 500


class RunTimeout(Exception):
    """--max-runtime 초과."""


def _now() -> str:
    return datetime.now(timezone.utc).strftime("%H:%M:%S")


def log(msg: str) -> None:
    print(f"[{_now()}] {msg}", flush=True)


def truncate(value: object, limit: int = PRINT_TRUNCATE) -> str:
    text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
    return text if len(text) <= limit else text[:limit] + f"... (+{len(text) - limit} chars)"


class AgentRunner:
    def __init__(self, args: argparse.Namespace) -> None:
        self.api_base = args.api_base.rstrip("/")
        self.yaml_path = Path(args.yaml_path)
        self.user_id = args.user_id
        self.workspace_id = args.workspace_id
        self.workspace_name = args.workspace_name
        self.user_input = args.user_input
        self.start_stage_index = args.start_stage_index
        self.read_timeout = args.read_timeout
        self.max_reconnects = args.max_reconnects
        self.max_runtime = args.max_runtime
        self.output_dir = Path(args.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.agent_name: str = ""
        self.session_id: str = ""
        self.final_status: str = "unknown"
        self.final_error: str = ""
        self.event_counts: dict[str, int] = {}
        self.stages_confirmed = 0
        self.started_at = time.monotonic()
        self.deadline = self.started_at + self.max_runtime

        # 백엔드 전역 AuthGate 통과용 서비스 토큰 (모든 데이터 라우트에 필수).
        # X-Admin-Token 은 어느 user_id 로든 대리 호출을 허용한다. 값은 환경변수로만
        # 받는다(코드/로그에 노출 금지) — 로컬은 .env, CI 는 Gitea secret.
        admin_token = os.environ.get("ADMIN_API_TOKEN", "")
        default_headers = {"X-Admin-Token": admin_token} if admin_token else {}
        if not admin_token:
            log("WARNING: ADMIN_API_TOKEN 미설정 — 인증이 필요한 백엔드는 401 로 거부됩니다")

        self.sse_client = httpx.Client(timeout=httpx.Timeout(30.0), headers=default_headers)
        # action/cancel 등 단발 요청은 SSE 스트림과 분리된 클라이언트로 보낸다
        self.api_client = httpx.Client(timeout=60.0, headers=default_headers)

        self.events_file = (self.output_dir / "events.jsonl").open("a", encoding="utf-8")

    def _remaining(self) -> float:
        return self.deadline - time.monotonic()

    # ------------------------------------------------------------- workspace

    def resolve_workspace(self) -> None:
        """workspace_id 를 확정한다.

        - workspace_id 가 주어지면(override) 그대로 사용.
        - 아니면 workspace_name 을 GET /workspaces/lookup 으로 UUID 해석.
        둘 다 없으면 SystemExit.
        """
        if self.workspace_id:
            log(f"Using explicit workspace_id: {self.workspace_id}")
            return
        if not self.workspace_name:
            raise SystemExit("ERROR: workspace_name 또는 workspace_id 중 하나를 지정해야 합니다.")

        log(f"Resolving workspace by name: {self.workspace_name!r} (user_id={self.user_id})")
        try:
            resp = self.api_client.get(
                f"{self.api_base}/workspaces/lookup",
                params={"user_id": self.user_id},
            )
        except httpx.HTTPError as exc:
            raise SystemExit(f"ERROR: workspaces/lookup 요청 실패: {exc}")
        if resp.status_code != 200:
            raise SystemExit(f"ERROR: workspaces/lookup HTTP {resp.status_code}: {resp.text[:500]}")
        try:
            items = resp.json().get("workspaces", [])
        except json.JSONDecodeError:
            raise SystemExit(f"ERROR: workspaces/lookup 응답이 JSON 이 아님: {resp.text[:300]}")

        def norm(s: object) -> str:
            return unicodedata.normalize("NFC", (s if isinstance(s, str) else "").strip())

        def field(w: dict, *keys: str) -> str:
            for k in keys:
                v = w.get(k)
                if v:
                    return v
            return ""

        target = norm(self.workspace_name)
        names, exact, ci = [], [], []
        for w in items:
            name = field(w, "name", "workspace_name", "title")
            wid = field(w, "workspace_id", "id")
            names.append(name)
            if not wid:
                continue
            if norm(name) == target:
                exact.append(wid)
            elif norm(name).lower() == target.lower():
                ci.append(wid)

        chosen = exact if exact else ci
        if len(chosen) == 1:
            self.workspace_id = chosen[0]
            log(f"Resolved workspace {self.workspace_name!r} -> {self.workspace_id}")
        elif len(chosen) > 1:
            raise SystemExit(
                f"ERROR: 워크스페이스 이름 {self.workspace_name!r} 이 여러 개 매칭됩니다: {chosen}"
            )
        else:
            available = ", ".join(repr(n) for n in names) or "(없음)"
            raise SystemExit(
                f"ERROR: {self.workspace_name!r} 에 해당하는 워크스페이스를 찾지 못했습니다. "
                f"사용 가능: {available}"
            )

    # ------------------------------------------------------------------ upload

    def upload(self) -> None:
        if not self.yaml_path.is_file():
            raise SystemExit(f"ERROR: YAML not found: {self.yaml_path}")

        log(f"Uploading agent YAML: {self.yaml_path}")
        with self.yaml_path.open("rb") as f:
            resp = self.api_client.post(
                f"{self.api_base}/upload-agent",
                params={"user_id": self.user_id},
                files={"file": (self.yaml_path.name, f, "application/x-yaml")},
            )
        if resp.status_code != 200:
            raise SystemExit(f"ERROR: upload-agent failed HTTP {resp.status_code}: {resp.text[:1000]}")

        data = resp.json()
        if not data.get("success") or not data.get("agent_name"):
            raise SystemExit(f"ERROR: upload-agent rejected: {json.dumps(data, ensure_ascii=False)[:1000]}")

        self.agent_name = data["agent_name"]
        log(f"Agent registered: {self.agent_name} (stages={data.get('stages')})")

    # ------------------------------------------------------------------ actions

    def cancel_session(self) -> None:
        if not (self.agent_name and self.session_id):
            return
        try:
            self.api_client.post(
                f"{self.api_base}/sse/agent/{quote(self.agent_name, safe='')}"
                f"/cancel/{self.session_id}"
            )
            log(f"Cancel requested for session {self.session_id}")
        except httpx.HTTPError as exc:
            log(f"WARNING: cancel request failed: {exc}")

    def _send_action(self, action: str, attempts: int = 5) -> bool:
        """Send a stage action; retry transient failures inside the server's 600s wait window."""
        url = (
            f"{self.api_base}/sse/agent/{quote(self.agent_name, safe='')}"
            f"/action/{self.session_id}"
        )
        for i in range(1, attempts + 1):
            try:
                resp = self.api_client.post(url, json={"action": action})
                if resp.status_code == 200:
                    log(f"  -> action '{action}' sent")
                    return True
                if resp.status_code in (400, 404):
                    # 세션 소멸/거부 — 재시도 무의미
                    log(f"  -> action '{action}' failed HTTP {resp.status_code}: {resp.text[:300]}")
                    return False
                log(f"  -> action '{action}' attempt {i}/{attempts} HTTP {resp.status_code}: {resp.text[:300]}")
            except httpx.HTTPError as exc:
                log(f"  -> action '{action}' attempt {i}/{attempts} request error: {exc}")
            if i < attempts:
                time.sleep(min(2 ** i, 60))
        return False

    # ------------------------------------------------------------------ events

    def _record(self, event: dict) -> None:
        event["_received_at"] = datetime.now(timezone.utc).isoformat()
        self.events_file.write(json.dumps(event, ensure_ascii=False) + "\n")
        self.events_file.flush()
        etype = event.get("type", "unknown")
        self.event_counts[etype] = self.event_counts.get(etype, 0) + 1

    def _handle(self, event: dict) -> bool:
        """Returns True when execution reached a terminal state."""
        self._record(event)
        etype = event.get("type", "unknown")

        if etype == "session_started":
            self.session_id = event.get("session_id", self.session_id)
            log(
                f"Session started: {self.session_id} "
                f"(total_stages={event.get('total_stages')}, start_index={event.get('start_stage_index')})"
            )
        elif etype == "stage_start":
            log(f"=== Stage start: index={event.get('stage_index')} {event.get('stage_name', '')}")
        elif etype == "task_start":
            log(f"  task start: {event.get('task_name', '?')}")
        elif etype == "task_iteration":
            log(f"  task iter : {event.get('task_name', '?')} #{event.get('iteration', '?')}")
        elif etype == "task_complete":
            log(f"  task done : {event.get('task_name', '?')}")
        elif etype == "stage_complete":
            log(f"=== Stage complete: index={event.get('stage_index')} — auto-confirming")
            if self._send_action("confirm"):
                self.stages_confirmed += 1
            else:
                # confirm 미전달 상태로 방치하면 서버 600s action-timeout 까지 세션이 잠긴다
                self.final_status = "error"
                self.final_error = "confirm action could not be delivered — cancelling session"
                log(f"ERROR: {self.final_error}")
                self.cancel_session()
                return True
        elif etype == "stage_error":
            # 서버는 stage_error 후 클라이언트 개입 없이는 같은 stage 를 재실행하거나
            # 실패 stage 를 건너뛰므로, CI 에서는 즉시 세션을 취소하고 실패 처리한다.
            self.final_status = "stage_error"
            self.final_error = str(event.get("error", ""))
            log(f"=== Stage ERROR: index={event.get('stage_index')} — {truncate(self.final_error)}")
            self.cancel_session()
            return True
        elif etype == "execution_complete":
            self.final_status = "completed"
            # 현행 백엔드는 final_outputs(stage 별 dict), 구버전 문서는 final_output
            final_output = event.get("final_outputs") or event.get("final_output", "")
            if not final_output:
                log("WARNING: execution_complete has neither final_outputs nor final_output")
            out_path = self.output_dir / "final_output.txt"
            out_path.write_text(
                final_output if isinstance(final_output, str)
                else json.dumps(final_output, ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
            log(f"Execution complete. final_output saved to {out_path}")
            return True
        elif etype == "execution_stopped":
            self.final_status = "stopped"
            self.final_error = str(event.get("reason") or event.get("message") or "")
            log(f"Execution STOPPED: {self.final_error}")
            return True
        elif etype == "error":
            self.final_status = "error"
            self.final_error = str(event.get("error", ""))
            log(f"Execution ERROR: {truncate(self.final_error)}")
            return True
        else:
            log(f"  event [{etype}]: {truncate({k: v for k, v in event.items() if k != 'type'})}")

        return False

    # ------------------------------------------------------------------ stream

    def _stream(self, mode: str) -> bool:
        """Open one SSE connection and consume events. Returns True on terminal event."""
        remaining = self._remaining()
        if remaining <= 0:
            raise RunTimeout()

        if mode == "start":
            url = f"{self.api_base}/sse/agent/{quote(self.agent_name, safe='')}/start"
            body = {
                "user_id": self.user_id,
                "workspace_id": self.workspace_id,
                "user_input": self.user_input,
                "start_stage_index": self.start_stage_index,
            }
        else:
            url = (
                f"{self.api_base}/sse/agent/{quote(self.agent_name, safe='')}"
                f"/reconnect/{self.session_id}"
            )
            body = {}

        # read timeout: 이벤트 간 무응답 상한이며 남은 실행 예산을 넘지 않게 잡는다
        timeout = httpx.Timeout(
            connect=30.0,
            read=min(self.read_timeout, max(remaining, 30.0)),
            write=30.0,
            pool=30.0,
        )

        log(f"SSE {mode}: {url}")
        with connect_sse(self.sse_client, "POST", url, json=body, timeout=timeout) as event_source:
            for sse in event_source.iter_sse():
                if self._remaining() <= 0:
                    raise RunTimeout()
                if not sse.data:
                    continue
                try:
                    event = json.loads(sse.data)
                except json.JSONDecodeError:
                    log(f"  (unparseable SSE data) {truncate(sse.data)}")
                    continue
                if self._handle(event):
                    return True
        return False

    def run(self) -> int:
        try:
            if not self.user_id:
                raise SystemExit("ERROR: user_id 를 지정해야 합니다 (--user-id).")
            self.resolve_workspace()
            self.upload()

            mode = "start"
            reconnects = 0
            while True:
                try:
                    if self._stream(mode):
                        break
                    # 스트림이 종료 이벤트 없이 닫힘 → 재연결 시도
                    raise ConnectionError("SSE stream ended without a terminal event")
                except RunTimeout:
                    self.final_status = "timeout"
                    self.final_error = (
                        f"--max-runtime {int(self.max_runtime)}s exceeded — cancelling session"
                    )
                    log(f"ERROR: {self.final_error}")
                    self.cancel_session()
                    break
                except SSEError as exc:
                    # 서버가 SSE 가 아닌 응답을 반환 (404/500 등)
                    self.final_status = "error"
                    self.final_error = f"SSE handshake failed ({mode}): {exc}"
                    log(f"ERROR: {self.final_error}")
                    break
                except (httpx.HTTPError, ConnectionError) as exc:
                    if self.final_status != "unknown":
                        break
                    reconnects += 1
                    if not self.session_id or reconnects > self.max_reconnects:
                        self.final_status = "error"
                        self.final_error = (
                            f"connection lost ({exc}); reconnect attempts exhausted "
                            f"({reconnects - 1}/{self.max_reconnects}). "
                            f"과거 세션 조회: GET {self.api_base}/history/sessions"
                            f"?user_id={self.user_id}&agent_name={self.agent_name}"
                        )
                        log(f"ERROR: {self.final_error}")
                        self.cancel_session()
                        break
                    wait = min(10 * reconnects, 60)
                    log(f"Connection lost ({exc}); reconnect {reconnects}/{self.max_reconnects} in {wait}s")
                    time.sleep(wait)
                    mode = "reconnect"
        except KeyboardInterrupt:
            # SIGINT / SIGTERM(핸들러가 KeyboardInterrupt 로 변환) — 세션 정리 후 종료
            self.final_status = "interrupted"
            self.final_error = "interrupted by SIGINT/SIGTERM — session cancelled"
            log(self.final_error)
            self.cancel_session()
        except SystemExit as exc:
            # upload() 의 ERROR 경로 — summary 에 남기고 실패 처리
            self.final_status = "error"
            self.final_error = str(exc)
            log(self.final_error)
        except Exception as exc:  # noqa: BLE001 — 어떤 실패든 summary 를 남긴다
            self.final_status = "error"
            self.final_error = f"unexpected {type(exc).__name__}: {exc}"
            log(f"ERROR: {self.final_error}")
        finally:
            try:
                self._write_summary()
            except Exception as exc:  # noqa: BLE001
                log(f"WARNING: summary write failed: {exc}")
            self.events_file.close()

        return 0 if self.final_status == "completed" else 1

    # ------------------------------------------------------------------ report

    def _write_summary(self) -> None:
        elapsed = int(time.monotonic() - self.started_at)
        status_icon = "✅" if self.final_status == "completed" else "❌"
        lines = [
            "# Agent Run Summary",
            "",
            "| 항목 | 값 |",
            "|------|-----|",
            f"| Agent | `{self.agent_name or '-'}` |",
            f"| YAML | `{self.yaml_path}` |",
            f"| Session | `{self.session_id or '-'}` |",
            f"| workspace | `{self.workspace_name or '-'}` (`{self.workspace_id or '-'}`) |",
            f"| user_id | `{self.user_id}` |",
            f"| 최종 상태 | {status_icon} `{self.final_status}` |",
            f"| 소요 시간 | {elapsed // 60}m {elapsed % 60}s (상한 {int(self.max_runtime) // 60}m) |",
            f"| 확인(confirm)한 stage 수 | {self.stages_confirmed} |",
        ]
        if self.final_error:
            lines += ["", f"**오류/중단 사유:** {self.final_error}"]
        lines += ["", "## 이벤트 수신 통계", "", "| type | count |", "|------|-------|"]
        for etype, count in sorted(self.event_counts.items()):
            lines.append(f"| {etype} | {count} |")
        summary = "\n".join(lines) + "\n"

        (self.output_dir / "summary.md").write_text(summary, encoding="utf-8")
        step_summary = os.environ.get("GITHUB_STEP_SUMMARY")
        if step_summary:
            with open(step_summary, "a", encoding="utf-8") as f:
                f.write(summary)
        log(f"Summary written ({self.final_status}, {elapsed}s)")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--yaml-path", required=True, help="실행할 Agent YAML 경로 (repo 상대/절대)")
    p.add_argument("--api-base", default="https://legalpoc.eroomai.com/api")
    p.add_argument("--user-id", default="", help="AgentBackend user_id (필수)")
    p.add_argument("--workspace-name", default="",
                   help="작업실 이름. GET /workspaces/lookup 으로 UUID 해석 (workspace-id 미지정 시)")
    p.add_argument("--workspace-id", default="",
                   help="작업실 UUID 직접 지정(override). 지정 시 workspace-name 무시")
    p.add_argument("--user-input", default="")
    p.add_argument("--start-stage-index", type=int, default=0)
    p.add_argument("--output-dir", default="agent_run_output")
    p.add_argument("--read-timeout", type=float, default=1800.0,
                   help="SSE 이벤트 간 최대 대기 초 (초과 시 재연결)")
    p.add_argument("--max-reconnects", type=int, default=5)
    p.add_argument("--max-runtime", type=float, default=9000.0,
                   help="총 실행 시간 상한 초. 초과 시 세션 cancel 후 exit 1. "
                        "Gitea/act_runner 의 3h 태스크 상한 및 job timeout-minutes 보다 짧게 잡을 것")
    return p.parse_args()


def _raise_interrupt(signum, frame):  # noqa: ARG001
    raise KeyboardInterrupt


def main() -> None:
    args = parse_args()
    # job timeout/취소 시 SIGTERM 이 오므로 KeyboardInterrupt 경로로 합류시켜 세션을 정리한다
    signal.signal(signal.SIGTERM, _raise_interrupt)
    runner = AgentRunner(args)
    sys.exit(runner.run())


if __name__ == "__main__":
    main()