# -*- coding: utf-8 -*-
"""SmartPaper R3.1 independent task executor.

Bootstrap is intentionally lightweight: the child process validates the local
request, claims the D1 task, writes a local claim handshake, and only then
imports the heavy paper-generation modules.  A Streamlit rerun/browser event
therefore cannot be mistaken for successful executor startup, and import or
pipeline failures after claim can be finalized through the authoritative
Worker so quota is never left reserved by an unowned task.
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import re
import sys
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import urlsplit

import requests

from config import APP_CONFIG

_log_handlers = [logging.StreamHandler(sys.stdout)]
_executor_log_path = str(os.environ.get("SMARTPAPER_EXECUTOR_LOG_PATH") or "").strip()
if _executor_log_path:
    try:
        Path(_executor_log_path).parent.mkdir(parents=True, exist_ok=True)
        _log_handlers.append(logging.FileHandler(_executor_log_path, encoding="utf-8"))
    except Exception:
        pass
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=_log_handlers,
)


CLAIM_HANDSHAKE_FILENAME = "executor_claim.json"


def _safe(value, limit=500):
    text = str(value or "")
    text = re.sub(r"(?i)(api[_ -]?key|authorization|bearer|token)\s*[:=]\s*[^\s,;]+", r"\1=[redacted]", text)
    return text[:limit]


def _api_origin() -> str:
    value = str(os.environ.get("SMARTPAPER_AUTH_API") or "").strip().rstrip("/")
    parsed = urlsplit(value)
    if parsed.scheme != "https" or not parsed.netloc or parsed.path not in ("", "/") or parsed.query or parsed.fragment:
        raise RuntimeError("SMARTPAPER_AUTH_API_invalid")
    return f"https://{parsed.netloc}"


def _service_key() -> str:
    value = str(os.environ.get("SMARTPAPER_SERVICE_KEY") or "").strip()
    if not value:
        raise RuntimeError("SMARTPAPER_SERVICE_KEY_missing")
    return value


def _atomic_json(path: Path, value: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, ensure_ascii=False), encoding="utf-8")
    os.replace(temporary, path)


def _claim_marker_path(task_dir: Path) -> Path:
    return task_dir / CLAIM_HANDSHAKE_FILENAME


class ExecutorApiError(RuntimeError):
    def __init__(self, error: str, status_code: int = 0, payload: dict | None = None):
        super().__init__(str(error or "executor_api_error"))
        self.error = str(error or "executor_api_error")
        self.status_code = int(status_code or 0)
        self.payload = dict(payload or {})


def _is_recoverable_generation_failure(exc: Exception) -> bool:
    """Classify only post-claim execution failures.

    This deliberately avoids importing the generation pipeline at module import
    time.  DeepSeek exception class names are stable public names inside this
    project and are inspected only after claim/bootstrap has succeeded.
    """
    if isinstance(exc, ExecutorApiError):
        return exc.status_code == 429 or exc.status_code >= 500
    class_name = exc.__class__.__name__
    if class_name in {"DeepSeekAuthenticationError", "DeepSeekBalanceError", "DeepSeekRequestError"}:
        return False
    if isinstance(exc, (requests.exceptions.Timeout, requests.exceptions.ConnectionError)):
        return True
    if class_name == "DeepSeekAPIError":
        return True
    text = str(exc or "").lower()
    transient_markers = (
        "connection", "timed out", "timeout", "temporarily unavailable",
        "streaming failed after", "api failed after", "stream was interrupted",
        "insufficient eligible references",
    )
    return any(marker in text for marker in transient_markers)


class ExecutorApi:
    def __init__(self):
        self.origin = _api_origin()
        self.headers = {
            "Authorization": f"Bearer {_service_key()}",
            "Accept": "application/json",
            "Content-Type": "application/json",
            "User-Agent": "SmartPaper-R3.1-Executor/1.0",
        }

    def post(self, path: str, payload: dict, *, timeout=20) -> dict:
        response = requests.post(
            self.origin + path,
            headers=self.headers,
            json=payload,
            timeout=timeout,
        )
        try:
            data = response.json()
        except Exception as exc:
            raise RuntimeError(f"executor_worker_non_json_http_{response.status_code}") from exc
        if not response.ok or data.get("ok") is False:
            error = str(data.get("error") or data.get("reason") or f"HTTP_{response.status_code}")
            raise ExecutorApiError(error, response.status_code, data)
        return data

    def claim(self, *, task_id: str, user_id: int, request_hash: str) -> dict:
        return self.post(
            "/executor/task/claim",
            {"task_id": task_id, "user_id": user_id, "request_hash": request_hash},
        )

    def heartbeat(self, payload: dict) -> dict:
        return self.post("/executor/task/heartbeat", payload)

    def finish(self, payload: dict) -> dict:
        return self.post("/executor/task/finish", payload, timeout=30)


class ExecutionState:
    def __init__(self, *, task_id: str, user_id: int, epoch: int, instance_id: str, api: ExecutorApi):
        self.task_id = task_id
        self.user_id = int(user_id)
        self.epoch = int(epoch)
        self.instance_id = str(instance_id)
        self.api = api
        self._lock = threading.Lock()
        self._last_activity = time.monotonic()
        self._progress_pct = 0
        self._current_node = "executor_claimed"
        self._checkpoint_seq = 0
        self._checkpoint_stage = ""
        self._ds_calls = 0
        self._s2_calls = 0
        self._oa_calls = 0
        self._cancel = False
        self._stop = threading.Event()
        self._thread = None

    def touch(self, *, progress_pct=None, current_node=None, checkpoint_seq=None, checkpoint_stage=None, ds_calls=None, s2_calls=None, oa_calls=None, activity_only=False):
        with self._lock:
            self._last_activity = time.monotonic()
            if not activity_only:
                if progress_pct is not None:
                    self._progress_pct = max(self._progress_pct, max(0, min(100, int(progress_pct))))
                if current_node:
                    self._current_node = str(current_node)[:240]
                if checkpoint_seq is not None and int(checkpoint_seq) >= self._checkpoint_seq:
                    self._checkpoint_seq = int(checkpoint_seq)
                    self._checkpoint_stage = str(checkpoint_stage or self._checkpoint_stage)[:120]
                if ds_calls is not None:
                    self._ds_calls = max(self._ds_calls, int(ds_calls))
                if s2_calls is not None:
                    self._s2_calls = max(self._s2_calls, int(s2_calls))
                if oa_calls is not None:
                    self._oa_calls = max(self._oa_calls, int(oa_calls))

    def cancelled(self) -> bool:
        with self._lock:
            return bool(self._cancel)

    def start_heartbeat(self):
        self._thread = threading.Thread(target=self._heartbeat_loop, name="smartpaper-executor-heartbeat", daemon=True)
        self._thread.start()

    def stop_heartbeat(self):
        self._stop.set()
        if self._thread and self._thread.is_alive():
            self._thread.join(timeout=2)

    def _payload(self):
        with self._lock:
            return {
                "task_id": self.task_id,
                "user_id": self.user_id,
                "execution_epoch": self.epoch,
                "executor_instance_id": self.instance_id,
                "current_node": self._current_node,
                "progress_pct": self._progress_pct,
                "checkpoint_seq": self._checkpoint_seq,
                "checkpoint_stage": self._checkpoint_stage,
                "ds_calls": self._ds_calls,
                "s2_calls": self._s2_calls,
                "oa_calls": self._oa_calls,
            }, time.monotonic() - self._last_activity

    def heartbeat_once(self):
        payload, _ = self._payload()
        result = self.api.heartbeat(payload)
        if result.get("cancel_requested"):
            with self._lock:
                self._cancel = True
        return result

    def _heartbeat_loop(self):
        interval = int(APP_CONFIG.get("executor_heartbeat_interval_seconds", 15))
        grace = int(APP_CONFIG.get("executor_activity_grace_seconds", 45))
        while not self._stop.wait(interval):
            payload, idle_seconds = self._payload()
            if idle_seconds > grace:
                logging.warning(
                    "[SP_EXECUTOR] heartbeat_paused task_id=%s idle_s=%.1f node=%s",
                    self.task_id, idle_seconds, payload.get("current_node"),
                )
                continue
            try:
                result = self.api.heartbeat(payload)
                if result.get("cancel_requested"):
                    with self._lock:
                        self._cancel = True
                    return
            except ExecutorApiError as exc:
                if exc.error in {
                    "executor_superseded", "generation_lock_lost",
                    "task_already_finalized", "task_cancel_requested",
                } or exc.status_code in {404, 409, 410}:
                    with self._lock:
                        self._cancel = True
                    logging.warning(
                        "[SP_EXECUTOR] executor_no_longer_authoritative task_id=%s error=%s",
                        self.task_id, _safe(exc.error, 180),
                    )
                    return
                logging.warning("[SP_EXECUTOR] heartbeat_error task_id=%s error=%s", self.task_id, _safe(exc, 180))
            except Exception as exc:
                logging.warning("[SP_EXECUTOR] heartbeat_error task_id=%s error=%s", self.task_id, _safe(exc, 180))


def _load_request(task_dir: Path) -> dict:
    path = task_dir / "request.json"
    if not path.is_file():
        raise RuntimeError("request_manifest_missing")
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise RuntimeError("request_manifest_invalid")
    return value


def _bootstrap_claim(*, task_dir: Path, task_id: str, user_id: int, request_data: dict, api: ExecutorApi) -> tuple[dict, int, str]:
    """Claim first and publish a PID-bound local handshake before heavy imports."""
    request_hash = str(request_data.get("request_hash") or "").strip()
    if not request_hash:
        raise RuntimeError("request_hash_missing")
    claim = api.claim(task_id=task_id, user_id=user_id, request_hash=request_hash)
    task = dict(claim.get("task") or {})
    execution = dict(claim.get("execution") or {})
    epoch = int(execution.get("execution_epoch") or task.get("execution_epoch") or 0)
    instance_id = str(execution.get("executor_instance_id") or "")
    if epoch <= 0 or not instance_id:
        raise RuntimeError("executor_claim_identity_missing")
    if int(task.get("target_words") or 0) != int(request_data.get("target_words") or 0):
        raise RuntimeError("target_words_claim_mismatch")
    try:
        _atomic_json(
            _claim_marker_path(task_dir),
            {
                "task_id": task_id,
                "user_id": int(user_id),
                "request_hash": request_hash,
                "pid": os.getpid(),
                "execution_epoch": epoch,
                "executor_instance_id": instance_id,
                "claimed_at": datetime.now(timezone.utc).isoformat(),
            },
        )
    except Exception as exc:
        try:
            api.finish(
                {
                    "task_id": task_id,
                    "user_id": int(user_id),
                    "execution_epoch": epoch,
                    "executor_instance_id": instance_id,
                    "status": "failed",
                    "safe_error_code": "E_EXECUTOR_HANDSHAKE_FAILED",
                    "safe_error_detail": _safe(exc, 800),
                    "current_node": "executor_handshake_failed",
                }
            )
        finally:
            raise
    return task, epoch, instance_id


def _finish_bootstrap_failure(api: ExecutorApi, *, task_id: str, user_id: int, epoch: int, instance_id: str, exc: Exception) -> None:
    try:
        api.finish(
            {
                "task_id": task_id,
                "user_id": user_id,
                "execution_epoch": epoch,
                "executor_instance_id": instance_id,
                "status": "failed",
                "safe_error_code": "E_EXECUTOR_BOOTSTRAP_FAILED",
                "safe_error_detail": _safe(exc, 800),
                "current_node": "executor_bootstrap_failed",
            }
        )
    except Exception as finish_exc:
        logging.warning("[SP_EXECUTOR] bootstrap_finalize_rejected task_id=%s error=%s", task_id, _safe(finish_exc, 220))


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--task-dir", required=True)
    parser.add_argument("--task-id", required=True)
    parser.add_argument("--user-id", required=True, type=int)
    args = parser.parse_args()

    task_dir = Path(args.task_dir).resolve()
    request_data = _load_request(task_dir)
    if str(request_data.get("task_id") or "") != args.task_id:
        raise RuntimeError("request_task_id_mismatch")
    if int(request_data.get("user_id") or 0) != int(args.user_id):
        raise RuntimeError("request_user_id_mismatch")

    api = ExecutorApi()
    task, epoch, instance_id = _bootstrap_claim(
        task_dir=task_dir,
        task_id=args.task_id,
        user_id=args.user_id,
        request_data=request_data,
        api=api,
    )
    state = ExecutionState(task_id=args.task_id, user_id=args.user_id, epoch=epoch, instance_id=instance_id, api=api)
    state.start_heartbeat()
    logging.info("[SP_EXECUTOR] claimed task_id=%s epoch=%s instance=%s", args.task_id, epoch, instance_id[:12])

    remote_checkpoint_seq = int(task.get("checkpoint_seq") or 0)
    local_checkpoints = list((task_dir / "checkpoints").glob("checkpoint_*.json")) if (task_dir / "checkpoints").is_dir() else []
    if remote_checkpoint_seq > 0 and not local_checkpoints:
        try:
            api.finish(
                {
                    "task_id": args.task_id,
                    "user_id": args.user_id,
                    "execution_epoch": epoch,
                    "executor_instance_id": instance_id,
                    "status": "failed",
                    "safe_error_code": "E_CHECKPOINT_MISSING",
                    "safe_error_detail": "Checkpoint metadata exists but local checkpoint files are unavailable after runtime replacement; quota released.",
                    "current_node": "checkpoint_missing",
                }
            )
        finally:
            state.stop_heartbeat()
        return 2

    try:
        # Heavy paper modules are deliberately imported only after D1 claim and
        # the local claim handshake. Any import/init failure is now observable
        # and can atomically release the reserved quota.
        from core.generation_runner import GenerationRunner

        runner = GenerationRunner(
            task_dir=task_dir,
            request_data=request_data,
            execution_epoch=epoch,
            progress_callback=state.touch,
            cancel_checker=state.cancelled,
        )
        result = runner.run()
        state.touch(progress_pct=100, current_node="package_validated")
        try:
            state.heartbeat_once()
        except Exception as exc:
            logging.warning("[SP_EXECUTOR] final_heartbeat_error task_id=%s error=%s", args.task_id, _safe(exc, 180))
        if state.cancelled():
            raise RuntimeError("task_cancel_requested")
        api.finish(
            {
                "task_id": args.task_id,
                "user_id": args.user_id,
                "execution_epoch": epoch,
                "executor_instance_id": instance_id,
                "status": "completed",
                "current_node": "package_validated",
                "output_filename": str(result.get("output_filename") or "")[:240],
            }
        )
        logging.info("[SP_EXECUTOR] completed task_id=%s epoch=%s", args.task_id, epoch)
        return 0
    except Exception as exc:
        if state.cancelled():
            logging.warning("[SP_EXECUTOR] stopped_non_authoritative task_id=%s epoch=%s", args.task_id, epoch)
            return 3
        if isinstance(exc, (ImportError, ModuleNotFoundError)):
            logging.exception("[SP_EXECUTOR] bootstrap_import_failed task_id=%s epoch=%s error=%s", args.task_id, epoch, _safe(exc, 240))
            _finish_bootstrap_failure(api, task_id=args.task_id, user_id=args.user_id, epoch=epoch, instance_id=instance_id, exc=exc)
            return 4
        if _is_recoverable_generation_failure(exc):
            logging.exception("[SP_EXECUTOR] recoverable_interruption task_id=%s epoch=%s error=%s", args.task_id, epoch, _safe(exc, 240))
            return 75
        logging.exception("[SP_EXECUTOR] failed task_id=%s epoch=%s error=%s", args.task_id, epoch, _safe(exc, 240))
        try:
            api.finish(
                {
                    "task_id": args.task_id,
                    "user_id": args.user_id,
                    "execution_epoch": epoch,
                    "executor_instance_id": instance_id,
                    "status": "failed",
                    "safe_error_code": "E_GENERATION_FAILED",
                    "safe_error_detail": _safe(exc, 800),
                    "current_node": "generation_failed",
                }
            )
        except Exception as finish_exc:
            logging.warning("[SP_EXECUTOR] failure_finalize_rejected task_id=%s error=%s", args.task_id, _safe(finish_exc, 220))
        return 1
    finally:
        state.stop_heartbeat()


if __name__ == "__main__":
    raise SystemExit(main())