"""自动任务排队。

大模型推理、量化因子、微调按这个顺序轮流补一批。
没有待领取的自动任务时立刻补下一种；距上一批已满 2 小时也会再补一批。
切片先不绑定节点，心跳到来时再改成 processing。高 IC 因子写入 alpha_factors。
"""

from __future__ import annotations

import random
import secrets
import uuid
from datetime import datetime, timedelta, timezone

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.clock import utcnow
from app.config import settings
from app.models import AlphaFactor, ComputeNode, ComputeTask, PushStatus, TaskStatus

ALPHA_BATCH_HOURS = 2
MIN_ABS_IC = 0.03
SLICE_COUNT = 3
QUEUE_HOLD_DAYS = 7
AUTO_CYCLE = ("llm_inference", "alpha_mining", "lora_training")
LLM_MODEL = "qwen2.5-3b-instruct-q4_k_m"
LORA_BASE = "qwen2.5-3b"
PROMPTS = (
    "用三句话说明去中心化算力网络做什么",
    "列出显卡推理时的三个资源限制",
    "解释量化因子的 IC 表示什么",
)
OPERATORS = (
    "ts_mean",
    "ts_std",
    "rank",
    "delta",
    "decay_linear",
    "correlation",
    "ts_rank",
    "scale",
)
SYMBOLS = ("BTCUSDT", "ETHUSDT", "BNBUSDT", "SOLUSDT")


def maybe_generate_auto_batch(
    session: Session,
    now: datetime | None = None,
    force: bool = False,
) -> list[ComputeTask]:
    """没有待领取的自动任务，或距上一批已满 2 小时，就按轮询生成下一种任务。"""
    moment = now or utcnow()
    pending, latest = _auto_queue_state(session)
    due = latest is None or moment - latest >= timedelta(hours=ALPHA_BATCH_HOURS)
    if pending and not due and not force:
        return []
    kind = _next_cycle_type(session)
    if kind == "llm_inference":
        return _create_llm_batch(session, moment)
    if kind == "lora_training":
        return _create_lora_batch(session, moment)
    return _create_alpha_batch(session, moment)


def maybe_generate_alpha_batch(
    session: Session,
    now: datetime | None = None,
    force: bool = False,
) -> list[ComputeTask]:
    return maybe_generate_auto_batch(session, now=now, force=force)


def assign_queued_auto_task(
    session: Session,
    node: ComputeNode,
    now: datetime | None = None,
) -> ComputeTask | None:
    """把最早的一条未绑定自动任务派给当前节点。节点已有执行中的自动任务时不再加派。"""
    moment = now or utcnow()
    holding = session.scalar(
        select(ComputeTask.id).where(
            ComputeTask.node_id == node.id,
            ComputeTask.task_type.in_(AUTO_CYCLE),
            ComputeTask.status == TaskStatus.PROCESSING,
        )
    )
    if holding is not None:
        return None
    task = session.scalar(
        select(ComputeTask)
        .where(
            ComputeTask.task_type.in_(AUTO_CYCLE),
            ComputeTask.status == TaskStatus.PENDING,
            ComputeTask.node_id.is_(None),
        )
        .order_by(ComputeTask.dispatched_at, ComputeTask.id)
        .limit(1)
    )
    if task is None:
        return None
    task.node_id = node.id
    task.status = TaskStatus.PROCESSING
    task.deadline_at = moment + timedelta(minutes=settings.task_deadline_minutes)
    session.flush()
    return task


def assign_queued_alpha_task(
    session: Session,
    node: ComputeNode,
    now: datetime | None = None,
) -> ComputeTask | None:
    return assign_queued_auto_task(session, node, now=now)


def ingest_alpha_factors(session: Session, task: ComputeTask) -> list[AlphaFactor]:
    """成功回传的挖掘结果里，只保存公式非空且绝对 IC 达到门槛的因子。"""
    if task.task_type != "alpha_mining" or task.status != TaskStatus.COMPLETED or not task.success:
        return []
    if task.node_id is None:
        return []
    payload = task.payload or {}
    saved: list[AlphaFactor] = []
    for item in _factor_items(task.result or {}):
        formula = str(item.get("formula") or item.get("alpha") or "").strip()
        ic = _as_float(item.get("ic"))
        if not formula or ic is None or abs(ic) < MIN_ABS_IC:
            continue
        symbol = item.get("symbol") or payload.get("symbol")
        seed = payload.get("seed")
        row = AlphaFactor(
            task_id=task.id,
            node_id=task.node_id,
            formula=formula,
            ic=ic,
            symbol=str(symbol) if symbol else None,
            seed=None if seed is None else str(seed),
            created_at=task.responded_at or utcnow(),
        )
        session.add(row)
        saved.append(row)
    if saved:
        session.flush()
    return saved


def _create_llm_batch(session: Session, moment: datetime) -> list[ComputeTask]:
    seed, rng, batch_id = _batch_identity()
    prompts = rng.sample(PROMPTS, k=SLICE_COUNT)
    tasks = [
        _queue_slice(
            session,
            moment,
            "llm_inference",
            {
                "auto": True,
                "batch_id": batch_id,
                "seed": seed,
                "model": LLM_MODEL,
                "prompt": prompt,
                "max_tokens": 256,
                "nonce": secrets.token_hex(16),
            },
        )
        for prompt in prompts
    ]
    session.flush()
    return tasks


def _create_lora_batch(session: Session, moment: datetime) -> list[ComputeTask]:
    seed, _rng, batch_id = _batch_identity()
    tasks = [
        _queue_slice(
            session,
            moment,
            "lora_training",
            {
                "auto": True,
                "batch_id": batch_id,
                "seed": seed,
                "base_model": LORA_BASE,
                "rank": rank,
                "steps": 100,
                "nonce": secrets.token_hex(16),
            },
        )
        for rank in (8, 16, 32)
    ]
    session.flush()
    return tasks


def _create_alpha_batch(session: Session, moment: datetime) -> list[ComputeTask]:
    seed, rng, batch_id = _batch_identity()
    combos = _operator_combos(rng, SLICE_COUNT)
    symbols = rng.sample(SYMBOLS, k=SLICE_COUNT)
    tasks = [
        _queue_slice(
            session,
            moment,
            "alpha_mining",
            {
                "auto": True,
                "batch_id": batch_id,
                "seed": seed,
                "operators": combo,
                "symbol": symbol,
                "nonce": secrets.token_hex(16),
            },
        )
        for combo, symbol in zip(combos, symbols)
    ]
    session.flush()
    return tasks


def _queue_slice(session: Session, moment: datetime, task_type: str, payload: dict) -> ComputeTask:
    task = ComputeTask(
        id=str(uuid.uuid4()),
        node_id=None,
        task_type=task_type,
        payload=payload,
        status=TaskStatus.PENDING,
        push_status=PushStatus.QUEUED,
        dispatched_at=moment,
        deadline_at=moment + timedelta(days=QUEUE_HOLD_DAYS),
    )
    session.add(task)
    return task


def _batch_identity() -> tuple[int, random.Random, str]:
    seed = secrets.randbelow(2**31)
    return seed, random.Random(seed), uuid.uuid4().hex[:12]


def _auto_queue_state(session: Session) -> tuple[int, datetime | None]:
    rows = session.scalars(
        select(ComputeTask).where(
            ComputeTask.task_type.in_(AUTO_CYCLE),
            ComputeTask.node_id.is_(None),
            ComputeTask.status == TaskStatus.PENDING,
        )
    ).all()
    pending = sum(1 for row in rows if (row.payload or {}).get("auto"))
    latest = _latest_auto_dispatched(session)
    return pending, latest


def _latest_auto_dispatched(session: Session) -> datetime | None:
    rows = session.scalars(
        select(ComputeTask)
        .where(ComputeTask.task_type.in_(AUTO_CYCLE))
        .order_by(ComputeTask.dispatched_at.desc(), ComputeTask.id.desc())
        .limit(100)
    ).all()
    for row in rows:
        if (row.payload or {}).get("auto"):
            return _as_utc(row.dispatched_at)
    return None


def _next_cycle_type(session: Session) -> str:
    rows = session.scalars(
        select(ComputeTask)
        .where(ComputeTask.task_type.in_(AUTO_CYCLE))
        .order_by(ComputeTask.dispatched_at.desc(), ComputeTask.id.desc())
        .limit(100)
    ).all()
    latest = None
    for row in rows:
        if (row.payload or {}).get("auto"):
            latest = row.task_type
            break
    if latest not in AUTO_CYCLE:
        return AUTO_CYCLE[0]
    return AUTO_CYCLE[(AUTO_CYCLE.index(latest) + 1) % len(AUTO_CYCLE)]


def _operator_combos(rng: random.Random, count: int) -> list[list[str]]:
    combos: list[list[str]] = []
    seen: set[tuple[str, ...]] = set()
    while len(combos) < count:
        size = 2 if rng.random() < 0.7 else 3
        combo = tuple(sorted(rng.sample(OPERATORS, k=size)))
        if combo in seen:
            continue
        seen.add(combo)
        combos.append(list(combo))
    return combos


def _factor_items(result: dict) -> list[dict]:
    factors = result.get("factors")
    if isinstance(factors, list):
        return [item for item in factors if isinstance(item, dict)]
    if any(key in result for key in ("formula", "alpha", "ic")):
        return [result]
    return []


def _as_utc(value: datetime | str | None) -> datetime | None:
    if value is None:
        return None
    parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) if isinstance(value, str) else value
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed


def _as_float(value: object) -> float | None:
    if isinstance(value, bool) or value is None:
        return None
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, str):
        try:
            return float(value)
        except ValueError:
            return None
    return None