"""Stop waiting Sunrise_pt-200-x1 validation tasks for selected accounts.

The worker processes accounts sequentially. For each account it discovers the
current target tasks, validates account/status/GPU locally, then stops them in
batches with that account's own xc-Token. The HTTP service remains available
for platform health and status probes after the worker finishes.
"""

import json
import os
import signal
import threading
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any

import requests

BASE_URL = os.environ.get("BASE_URL", "https://modelhub.org.cn").rstrip("/")
TASK_PAGE_ENDPOINT = "/api/adapt/task/page"
STOP_TASK_ENDPOINT = "/api/async/task/stop-create-contest-task"
TARGET_GPU = "Sunrise_pt-200-x1"
TARGET_STATUS = "waiting"
STRATEGY_ID = os.environ.get("STRATEGY_ID", "")

HTTP_HOST = "0.0.0.0"
HTTP_PORT = int(os.environ.get("PORT", "8080"))
PAGE_SIZE = 100
BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "50"))
MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "3"))
REQUEST_TIMEOUT = int(os.environ.get("REQUEST_TIMEOUT", "30"))

TARGETS = [
    {
        "account": "jiangxiaowen",
        "xc_token": "88d5fee9f1fe4f7583f11a9d3702dc85",
    },
    {
        "account": "l112233",
        "xc_token": "40cb6910dc9a442a816298a228da65ac",
    },
    {
        "account": "l11223344",
        "xc_token": "e1c0db2959e5411f9342c8550b03f6e9",
    },
]

_shutdown = threading.Event()
_state: dict[str, Any] = {
    "strategy_id": STRATEGY_ID,
    "target_gpu": TARGET_GPU,
    "target_status": TARGET_STATUS,
    "accounts": {
        target["account"]: {
            "phase": "pending",
            "total": 0,
            "stopped": 0,
            "failed": 0,
            "failed_task_ids": [],
        }
        for target in TARGETS
    },
    "phase": "starting",  # starting | stopping | done | partial_failure | error
    "total": 0,
    "stopped": 0,
    "failed": 0,
    "failed_task_ids": [],
    "started_at": None,
    "finished_at": None,
    "error": None,
}


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


class Handler(BaseHTTPRequestHandler):
    def do_GET(self) -> None:  # noqa: N802 - BaseHTTPRequestHandler API
        if self.path == "/health":
            self._json({"status": "ok"})
        elif self.path == "/status":
            self._json(_state)
        else:
            self._json({"error": "not found"}, 404)

    def _json(self, body: dict[str, Any], status: int = 200) -> None:
        payload = json.dumps(body, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def log_message(self, fmt: str, *args: object) -> None:
        print(f"[http] {self.address_string()} {fmt % args}", flush=True)


def _run_http() -> None:
    server = ThreadingHTTPServer((HTTP_HOST, HTTP_PORT), Handler)
    server.timeout = 1
    print(f"[http] listening on {HTTP_HOST}:{HTTP_PORT}", flush=True)
    while not _shutdown.is_set():
        server.handle_request()
    server.server_close()


def _fetch_target_task_ids(account: str, xc_token: str) -> list[int]:
    """Return only this account's waiting tasks on the exact target GPU."""
    headers = {"xc-Token": xc_token}
    task_ids: list[int] = []
    page = 1

    while True:
        response = requests.get(
            f"{BASE_URL}{TASK_PAGE_ENDPOINT}",
            headers=headers,
            params={
                "current": page,
                "pageSize": PAGE_SIZE,
                "status": TARGET_STATUS,
                "gpuType": TARGET_GPU,
            },
            timeout=REQUEST_TIMEOUT,
        )
        response.raise_for_status()
        result = response.json()
        if result.get("code") != 0:
            raise RuntimeError(
                f"{account}: task query failed: {result.get('message', 'unknown error')}"
            )

        data = result.get("data") or {}
        for record in data.get("records") or []:
            record_account = record.get("userAccount")
            record_status = str(record.get("status", "")).lower()
            record_gpu = record.get("gpuType")
            if (
                record_account != account
                or record_status != TARGET_STATUS
                or record_gpu != TARGET_GPU
            ):
                print(
                    f"[query] {account}: rejected mismatched record "
                    f"taskId={record.get('taskId')} account={record_account!r} "
                    f"status={record_status!r} gpuType={record_gpu!r}",
                    flush=True,
                )
                continue
            try:
                task_ids.append(int(record["taskId"]))
            except (KeyError, TypeError, ValueError):
                print(f"[query] {account}: rejected non-numeric task ID", flush=True)

        pages = int(data.get("pages") or 0)
        if page >= pages:
            break
        page += 1

    unique_task_ids = list(dict.fromkeys(task_ids))
    print(
        f"[query] {account}: found {len(unique_task_ids)} "
        f"{TARGET_STATUS} tasks on {TARGET_GPU}",
        flush=True,
    )
    return unique_task_ids


def _stop_batch(account: str, xc_token: str, task_ids: list[int]) -> bool:
    headers = {"Content-Type": "application/json", "xc-Token": xc_token}
    url = f"{BASE_URL}{STOP_TASK_ENDPOINT}"
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            response = requests.put(
                url,
                headers=headers,
                json={"taskIds": task_ids},
                timeout=REQUEST_TIMEOUT,
            )
            try:
                result = response.json()
            except ValueError:
                result = {"message": response.text[:500]}

            if response.ok and result.get("code") == 0:
                print(
                    f"[stop] {account}: stopped {len(task_ids)} task IDs",
                    flush=True,
                )
                return True
            print(
                f"[stop] {account}: attempt {attempt}/{MAX_RETRIES} failed "
                f"for {len(task_ids)} tasks: HTTP {response.status_code}, {result}",
                flush=True,
            )
        except requests.RequestException as exc:
            print(
                f"[stop] {account}: attempt {attempt}/{MAX_RETRIES} "
                f"request error: {exc}",
                flush=True,
            )

        if attempt < MAX_RETRIES and not _shutdown.wait(attempt):
            continue
        if _shutdown.is_set():
            break
    return False


def _run_worker() -> None:
    _state["started_at"] = _now()
    _state["phase"] = "stopping"
    try:
        for target in TARGETS:
            if _shutdown.is_set():
                break
            account = target["account"]
            account_state = _state["accounts"][account]
            account_state["phase"] = "querying"
            task_ids = _fetch_target_task_ids(account, target["xc_token"])
            account_state["total"] = len(task_ids)
            _state["total"] += len(task_ids)
            account_state["phase"] = "stopping"

            for start in range(0, len(task_ids), BATCH_SIZE):
                if _shutdown.is_set():
                    break
                batch = task_ids[start : start + BATCH_SIZE]
                if _stop_batch(account, target["xc_token"], batch):
                    _state["stopped"] += len(batch)
                    account_state["stopped"] += len(batch)
                else:
                    _state["failed"] += len(batch)
                    _state["failed_task_ids"].extend(batch)
                    account_state["failed"] += len(batch)
                    account_state["failed_task_ids"].extend(batch)

            account_state["phase"] = (
                "done" if account_state["failed"] == 0 else "partial_failure"
            )

        _state["phase"] = "done" if _state["failed"] == 0 else "partial_failure"
    except Exception as exc:  # exposed through /status for diagnosis
        _state["phase"] = "error"
        _state["error"] = str(exc)
        print(f"[stop] fatal error: {exc}", flush=True)
    finally:
        _state["finished_at"] = _now()
        print(f"[stop] completed: {_state}", flush=True)


def _handle_signal(signum: int, _frame: Any) -> None:
    print(f"[main] received signal {signum}; shutting down", flush=True)
    _shutdown.set()


def main() -> None:
    signal.signal(signal.SIGTERM, _handle_signal)
    signal.signal(signal.SIGINT, _handle_signal)

    http_thread = threading.Thread(target=_run_http, daemon=False)
    http_thread.start()
    threading.Thread(target=_run_worker, daemon=True).start()

    _shutdown.wait()
    http_thread.join(timeout=5)


if __name__ == "__main__":
    main()