"""随机任务下发与回传校验。

节点不暴露回调地址时，任务留在待拉取队列，客户端用
GET /api/v1/node/{sn}/tasks 领取。配置了 callback_url 时，调度器会额外 POST 一次。
回传必须带上任务里的 nonce，避免未收到任务的客户端直接上报成功。
"""

from __future__ import annotations

import hmac
import json
import logging
import random
import secrets
import uuid
from datetime import datetime, timedelta
from urllib.error import URLError
from urllib.request import Request, urlopen

from sqlalchemy import select
from sqlalchemy.orm import Session, selectinload

from app.clock import utcnow
from app.config import settings
from app.errors import AppError
from app.models import ComputeNode, ComputeTask, NodeStatus, PushStatus, TaskStatus
from app.services.alpha_queue import AUTO_CYCLE, ingest_alpha_factors
from app.services.common import get_node_by_sn

logger = logging.getLogger(__name__)

ADMIN_TASK_TYPES = {
    "alpha_mining": "量化因子挖掘",
    "llm_inference": "大模型推理",
    "lora_training": "QLoRA 微调",
    "stress_test": "简易压测",
}
_BUSY = (TaskStatus.PENDING, TaskStatus.PROCESSING)


def create_admin_tasks(
    session: Session,
    task_type: str,
    payload: dict | None,
    target_node_sn: str | None = None,
    now: datetime | None = None,
) -> list[ComputeTask]:
    """向指定节点或当前全部在线节点各放一条任务。payload 里补上一次性 nonce。"""
    if task_type not in ADMIN_TASK_TYPES:
        raise AppError("不支持的任务类型", 400)
    moment = now or utcnow()
    body = dict(payload or {})
    if len(json.dumps(body, ensure_ascii=False)) > 8000:
        raise AppError("任务参数过长", 400)
    targets = _dispatch_targets(session, target_node_sn, moment)
    deadline = moment + timedelta(minutes=settings.task_deadline_minutes)
    tasks: list[ComputeTask] = []
    for node in targets:
        task_payload = dict(body)
        task_payload["nonce"] = secrets.token_hex(16)
        task = ComputeTask(
            id=str(uuid.uuid4()),
            node_id=node.id,
            task_type=task_type,
            payload=task_payload,
            status=TaskStatus.PENDING,
            push_status=PushStatus.QUEUED,
            dispatched_at=moment,
            deadline_at=deadline,
        )
        session.add(task)
        tasks.append(task)
    session.flush()
    for node, task in zip(targets, tasks):
        _push_task(node, task)
    session.flush()
    return tasks


def list_admin_tasks(session: Session) -> list[ComputeTask]:
    return list(
        session.scalars(
            select(ComputeTask)
            .options(selectinload(ComputeTask.node))
            .order_by(ComputeTask.dispatched_at.desc(), ComputeTask.id.desc())
        ).all()
    )


def result_preview(result: dict | None) -> str:
    if not result:
        return ""
    for key in ("formula", "alpha", "output", "summary"):
        value = result.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()[:120]
    for key in ("patch_bytes", "patch_size", "adapter_bytes"):
        if result.get(key) is not None:
            return f"补丁 {result[key]}"
    text = json.dumps(result, ensure_ascii=False)
    return text if len(text) <= 120 else text[:117] + "..."


def _dispatch_targets(session: Session, target_node_sn: str | None, moment: datetime) -> list[ComputeNode]:
    sn = (target_node_sn or "").strip()
    if sn:
        node = get_node_by_sn(session, sn)
        if node.is_banned:
            raise AppError("节点已封禁，不能下发任务", 409)
        if node.license_expires_at <= moment:
            raise AppError("节点授权已到期，不能下发任务", 409)
        return [node]
    nodes = session.scalars(
        select(ComputeNode).where(
            ComputeNode.status == NodeStatus.ONLINE,
            ComputeNode.license_expires_at > moment,
            ComputeNode.is_banned.is_(False),
        )
    ).all()
    if not nodes:
        raise AppError("当前没有可下发的在线节点", 409)
    return list(nodes)


def dispatch_random_tasks(session: Session, now: datetime | None = None) -> list[ComputeTask]:
    """向当前在线、授权仍有效、且没有未完成任务的节点随机下发一轮任务。"""
    moment = now or utcnow()
    expire_overdue_tasks(session, moment)
    # 未绑定的自动切片不算占用节点。推理、因子和微调单独排队，不挡住模拟任务。
    pending_node_ids = set(
        session.scalars(
            select(ComputeTask.node_id).where(
                ComputeTask.node_id.is_not(None),
                ComputeTask.status.in_(_BUSY),
                ComputeTask.task_type.notin_(AUTO_CYCLE),
            )
        ).all()
    )
    candidates = [
        node
        for node in session.scalars(
            select(ComputeNode).where(
                ComputeNode.status == NodeStatus.ONLINE,
                ComputeNode.license_expires_at > moment,
            )
        ).all()
        if node.id not in pending_node_ids
    ]
    if not candidates:
        return []

    batch = min(len(candidates), settings.dispatch_batch_size)
    chosen = random.sample(candidates, batch)
    deadline = moment + timedelta(minutes=settings.task_deadline_minutes)
    tasks: list[ComputeTask] = []
    for node in chosen:
        task = ComputeTask(
            id=str(uuid.uuid4()),
            node_id=node.id,
            task_type="simulated_compute",
            payload={
                "op": "simulated_compute",
                "nonce": secrets.token_hex(16),
                "seed": secrets.token_hex(8),
            },
            status=TaskStatus.PENDING,
            push_status=PushStatus.QUEUED,
            dispatched_at=moment,
            deadline_at=deadline,
        )
        session.add(task)
        tasks.append(task)
    session.flush()
    for node, task in zip(chosen, tasks):
        _push_task(node, task)
    session.flush()
    return tasks


def list_pending_tasks(session: Session, sn: str, now: datetime | None = None) -> list[ComputeTask]:
    moment = now or utcnow()
    node = get_node_by_sn(session, sn.strip())
    tasks = session.scalars(
        select(ComputeTask)
        .where(ComputeTask.node_id == node.id, ComputeTask.status.in_(_BUSY))
        .order_by(ComputeTask.dispatched_at, ComputeTask.id)
    ).all()
    visible: list[ComputeTask] = []
    for task in tasks:
        if task.deadline_at < moment:
            task.status = TaskStatus.EXPIRED
            continue
        if task.status == TaskStatus.PENDING:
            task.status = TaskStatus.PROCESSING
        visible.append(task)
    session.flush()
    return visible


def report_task(
    session: Session,
    sn: str,
    task_id: str,
    success: bool,
    nonce: str,
    now: datetime | None = None,
    result: dict | None = None,
) -> ComputeTask:
    """校验任务归属、时限和 nonce。只有校验通过且 success 为真才记为有效响应。"""
    moment = now or utcnow()
    node = get_node_by_sn(session, sn.strip())
    task = session.get(ComputeTask, task_id)
    if task is None:
        raise AppError("任务不存在", 404)
    if task.node_id != node.id:
        raise AppError("任务不属于该节点", 403)
    if task.status not in _BUSY:
        raise AppError("任务已结束，不能重复回传", 409)
    if task.deadline_at < moment:
        task.status = TaskStatus.EXPIRED
        session.flush()
        raise AppError("任务已超过回传时限", 409)

    expected = str((task.payload or {}).get("nonce", ""))
    if not _nonce_matches(nonce, expected):
        raise AppError("任务响应校验失败")

    task.responded_at = moment
    task.success = success
    task.result = result or None
    task.status = TaskStatus.COMPLETED if success else TaskStatus.FAILED
    session.flush()
    if task.status == TaskStatus.COMPLETED and task.task_type == "alpha_mining":
        ingest_alpha_factors(session, task)
    return task


def expire_overdue_tasks(session: Session, now: datetime | None = None) -> int:
    moment = now or utcnow()
    tasks = session.scalars(
        select(ComputeTask).where(
            ComputeTask.node_id.is_not(None),
            ComputeTask.status.in_(_BUSY),
            ComputeTask.deadline_at < moment,
        )
    ).all()
    for task in tasks:
        task.status = TaskStatus.EXPIRED
    if tasks:
        session.flush()
    return len(tasks)


def _nonce_matches(provided: str, expected: str) -> bool:
    if not expected or len(provided) != len(expected):
        return False
    return hmac.compare_digest(provided, expected)


def _push_task(node: ComputeNode, task: ComputeTask) -> None:
    if not node.callback_url:
        task.push_status = PushStatus.QUEUED
        return
    body = json.dumps(
        {
            "task_id": task.id,
            "sn": node.sn,
            "task_type": task.task_type,
            "payload": task.payload,
            "deadline_at": task.deadline_at.isoformat(),
        }
    ).encode()
    request = Request(
        node.callback_url,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urlopen(request, timeout=2) as response:
            if 200 <= response.status < 300:
                task.push_status = PushStatus.PUSHED
                task.push_error = None
                return
            task.push_status = PushStatus.FAILED
            task.push_error = f"HTTP {response.status}"
    except (URLError, TimeoutError, OSError) as exc:
        task.push_status = PushStatus.FAILED
        task.push_error = str(exc)[:255]
        logger.warning("task push failed sn=%s task=%s error=%s", node.sn, task.id, exc)