# -*- coding: utf-8 -*-
"""调度任务状态持久化（跨进程重启 / 容器重建）。

每条任务一条 kv 记录 ``task:state:{name}``，字段见 :class:`TaskState`。
调度器只认落盘的 ``due_at`` / ``status`` / ``stat_date`` / ``progress``，
不再靠启动时把 ``next_run=0`` 猜「要不要立刻跑」。
"""

from __future__ import annotations

import json
import time
from dataclasses import asdict, dataclass
from datetime import timedelta
from typing import Any, Mapping

from .repository import Repository
from .timeutil import parse_hhmm, shanghai_date, shanghai_datetime

STATUS_PENDING = "pending"
STATUS_RUNNING = "running"
STATUS_DONE = "done"
STATUS_FAILED = "failed"
STATUS_DEFERRED = "deferred"

_ACTIVE = frozenset({STATUS_PENDING, STATUS_RUNNING, STATUS_FAILED, STATUS_DEFERRED})


@dataclass
class TaskState:
    name: str
    status: str
    due_at: float
    stat_date: str | None = None
    progress: dict[str, Any] | None = None
    updated_at: float = 0.0
    error: str | None = None

    def as_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        # JSON 友好：去掉 None progress 也可保留
        return payload


def state_key(name: str) -> str:
    return f"task:state:{name}"


def load_state(repo: Repository, name: str) -> TaskState | None:
    raw = repo.kv_get(state_key(name))
    if not raw:
        return None
    try:
        data = json.loads(raw)
    except (TypeError, ValueError, json.JSONDecodeError):
        return None
    if not isinstance(data, dict) or not data.get("name"):
        return None
    try:
        due_at = float(data.get("due_at") or 0)
        updated_at = float(data.get("updated_at") or 0)
    except (TypeError, ValueError):
        return None
    progress = data.get("progress")
    if progress is not None and not isinstance(progress, dict):
        progress = None
    return TaskState(
        name=str(data.get("name") or name),
        status=str(data.get("status") or STATUS_PENDING),
        due_at=due_at,
        stat_date=(str(data["stat_date"]) if data.get("stat_date") else None),
        progress=progress,
        updated_at=updated_at,
        error=(str(data["error"]) if data.get("error") else None),
    )


def save_state(repo: Repository, state: TaskState) -> None:
    repo.kv_set(state_key(state.name), json.dumps(state.as_dict(), ensure_ascii=False))


def _now(now: float | None) -> float:
    return time.time() if now is None else float(now)


def recover_stale_running(
    repo: Repository,
    name: str,
    *,
    now: float | None = None,
) -> TaskState | None:
    """进程已死但仍为 ``running``：改为 ``pending``，保留 progress / stat_date。

    ``due_at`` 取 ``min(原 due_at, now)``，使过期/中断任务可立即续跑，
    未到期的（极少见）仍尊重原 due_at。
    """
    current = _now(now)
    state = load_state(repo, name)
    if state is None or state.status != STATUS_RUNNING:
        return state
    state.status = STATUS_PENDING
    state.due_at = min(float(state.due_at or current), current)
    state.updated_at = current
    save_state(repo, state)
    return state


def mark_running(
    repo: Repository,
    name: str,
    *,
    stat_date: str | None,
    now: float | None = None,
    progress: Mapping[str, Any] | None = None,
    due_at: float | None = None,
) -> TaskState:
    current = _now(now)
    prev = load_state(repo, name)
    state = TaskState(
        name=name,
        status=STATUS_RUNNING,
        due_at=float(due_at if due_at is not None else (prev.due_at if prev else current)),
        stat_date=stat_date if stat_date is not None else (prev.stat_date if prev else None),
        progress=dict(progress) if progress is not None else (prev.progress if prev else None),
        updated_at=current,
        error=None,
    )
    save_state(repo, state)
    return state


def mark_done(
    repo: Repository,
    name: str,
    *,
    due_at: float,
    now: float | None = None,
    stat_date: str | None = None,
    progress: Mapping[str, Any] | None = None,
) -> TaskState:
    """成功：记 done，并写入下一次 ``due_at``（下一条墙钟/间隔）。"""
    current = _now(now)
    state = TaskState(
        name=name,
        status=STATUS_DONE,
        due_at=float(due_at),
        stat_date=stat_date,
        progress=dict(progress) if progress is not None else None,
        updated_at=current,
        error=None,
    )
    save_state(repo, state)
    return state


def mark_failed(
    repo: Repository,
    name: str,
    *,
    due_at: float,
    now: float | None = None,
    stat_date: str | None = None,
    progress: Mapping[str, Any] | None = None,
    error: str | None = None,
) -> TaskState:
    current = _now(now)
    prev = load_state(repo, name)
    state = TaskState(
        name=name,
        status=STATUS_FAILED,
        due_at=float(due_at),
        stat_date=stat_date if stat_date is not None else (prev.stat_date if prev else None),
        progress=(
            dict(progress) if progress is not None
            else (prev.progress if prev else None)
        ),
        updated_at=current,
        error=error,
    )
    save_state(repo, state)
    return state


def mark_deferred(
    repo: Repository,
    name: str,
    *,
    due_at: float,
    now: float | None = None,
    stat_date: str | None = None,
    progress: Mapping[str, Any] | None = None,
    error: str | None = None,
) -> TaskState:
    """主动延后（静默窗 / seed 避让等）：保留进度，写未来 due_at。"""
    current = _now(now)
    prev = load_state(repo, name)
    state = TaskState(
        name=name,
        status=STATUS_DEFERRED,
        due_at=float(due_at),
        stat_date=stat_date if stat_date is not None else (prev.stat_date if prev else None),
        progress=(
            dict(progress) if progress is not None
            else (prev.progress if prev else None)
        ),
        updated_at=current,
        error=error,
    )
    save_state(repo, state)
    return state


def ensure_pending(
    repo: Repository,
    name: str,
    *,
    due_at: float,
    stat_date: str | None,
    now: float | None = None,
    overwrite_done: bool = True,
) -> TaskState:
    """保证存在一条 pending（如 seed 成功后挂上当晚 seal）。

    同日未完成：取更早的 due_at。
    跨日：不得把新一天的 due 粘到昨天的 ``stat_date`` 上——
    若已有更早日期的未完成记录，保留它（跨天续跑优先）；
    若已有更新日期的未完成，则重置为调用方的 due/stat_date。
    ``overwrite_done`` 时覆盖已完成记录。
    """
    current = _now(now)
    prev = load_state(repo, name)
    if prev is not None and prev.status in _ACTIVE and prev.status != STATUS_RUNNING:
        if (
            stat_date is not None
            and prev.stat_date
            and prev.stat_date != stat_date
        ):
            # 未完成的更早归属日优先保留，避免 seed 挂今晚时冲掉昨天续跑。
            if prev.stat_date < stat_date:
                return prev
            # prev 归属日更新：按调用方重置，禁止跨日字段拼接。
            state = TaskState(
                name=name,
                status=STATUS_PENDING,
                due_at=float(due_at),
                stat_date=stat_date,
                progress=None,
                updated_at=current,
                error=None,
            )
            save_state(repo, state)
            return state
        # 同日（或缺一边 stat_date）：取更早 due_at，补上缺失的 stat_date
        due = min(float(prev.due_at or due_at), float(due_at))
        sd = prev.stat_date or stat_date
        if due == prev.due_at and sd == prev.stat_date and prev.status == STATUS_PENDING:
            return prev
        prev.due_at = due
        prev.stat_date = sd
        prev.status = STATUS_PENDING
        prev.updated_at = current
        save_state(repo, prev)
        return prev
    if prev is not None and prev.status == STATUS_RUNNING:
        return prev
    if prev is not None and prev.status == STATUS_DONE and not overwrite_done:
        return prev
    state = TaskState(
        name=name,
        status=STATUS_PENDING,
        due_at=float(due_at),
        stat_date=stat_date,
        progress=None,
        updated_at=current,
        error=None,
    )
    save_state(repo, state)
    return state


def bootstrap_from_last(
    repo: Repository,
    name: str,
    *,
    wall_due_at: float,
    now: float | None = None,
    retry_seconds: float = 300.0,
) -> TaskState | None:
    """升级/冷启动：无 ``task:state`` 时从 ``task:last`` 回填。

    - 成功 → ``done`` + ``due_at=wall_due_at``
    - 失败/跳过 → ``failed``/``deferred``，``due_at`` 为补跑点（已过则立刻）
    """
    current = _now(now)
    existing = load_state(repo, name)
    if existing is not None:
        return existing
    raw = repo.kv_get(f"task:last:{name}")
    if not raw:
        return None
    try:
        data = json.loads(raw)
    except (TypeError, ValueError, json.JSONDecodeError):
        return None
    if not isinstance(data, dict):
        return None
    ok = bool(data.get("ok"))
    skipped = bool(data.get("skipped"))
    try:
        finished = float(data.get("finished_at") or data.get("started_at") or current)
    except (TypeError, ValueError):
        finished = current
    summary = data.get("summary")
    stat_date = None
    if isinstance(summary, dict) and summary.get("stat_date"):
        stat_date = str(summary["stat_date"])
    error = data.get("error")
    error_s = str(error) if error else None
    if ok and not skipped:
        state = TaskState(
            name=name,
            status=STATUS_DONE,
            due_at=float(wall_due_at),
            stat_date=stat_date,
            updated_at=current,
            error=None,
        )
        save_state(repo, state)
        return state
    due = finished + float(retry_seconds)
    if due > current:
        pass
    else:
        due = current
    state = TaskState(
        name=name,
        status=STATUS_DEFERRED if skipped else STATUS_FAILED,
        due_at=float(due),
        stat_date=stat_date,
        updated_at=current,
        error=error_s,
    )
    save_state(repo, state)
    return state


def bootstrap_seed_markers(
    repo: Repository,
    name: str,
    *,
    slot_hhmm: str,
    wall_due_at: float,
    now: float | None = None,
) -> TaskState | None:
    """无 state 时：若存在未 seal 的 seed 壳，挂 pending（跨天优先昨天）。

    已过当日 ``slot_hhmm``（或归属日早于今天）→ ``due_at=now`` 尽快补跑。
    """
    del wall_due_at  # 未完成 seed 不走纯墙钟
    current = _now(now)
    existing = load_state(repo, name)
    if existing is not None:
        return existing
    if name not in {"snapshot", "fans"}:
        return None
    today = shanghai_date(current)
    yday = (shanghai_datetime(current) - timedelta(days=1)).strftime("%Y-%m-%d")
    hour, minute = parse_hhmm(slot_hhmm)
    today_slot = shanghai_datetime(current).replace(
        hour=hour, minute=minute, second=0, microsecond=0,
    ).timestamp()

    def _seeded_unfinished(day: str) -> bool:
        if name == "snapshot":
            if repo.kv_get(f"snapshot:seeded:{day}") != "1":
                return False
            return repo.kv_get(f"snapshot:completed:{day}") != "1"
        # fans：seed 键在成功 seal 后删除，有键即未完成
        return repo.kv_get(f"fans:seeded:{day}") == "1"

    for day in (yday, today):
        if not _seeded_unfinished(day):
            continue
        if day < today or current >= today_slot:
            due = current
        else:
            due = float(today_slot)
        return ensure_pending(
            repo, name, due_at=due, stat_date=day, now=current,
        )
    return None


def startup_next_run(
    repo: Repository,
    name: str,
    *,
    wall_due_at: float,
    now: float | None = None,
    slot_hhmm: str | None = None,
    retry_seconds: float = 300.0,
) -> float:
    """常驻启动时的 ``ScheduledTask.next_run``：只认落盘状态。

    - 无状态时先 ``bootstrap_from_last`` / seed 标记回填
    - 恢复残留 ``running`` → ``pending``
    - ``due_at`` 是唯一时间门闩（失败延时不会被提前）
    """
    current = _now(now)
    if load_state(repo, name) is None:
        bootstrap_from_last(
            repo, name, wall_due_at=wall_due_at, now=current,
            retry_seconds=retry_seconds,
        )
        if slot_hhmm:
            bootstrap_seed_markers(
                repo, name, slot_hhmm=slot_hhmm,
                wall_due_at=wall_due_at, now=current,
            )
    state = recover_stale_running(repo, name, now=current)
    if state is None:
        return float(wall_due_at)
    due = float(state.due_at or 0)
    if due <= 0:
        return float(wall_due_at)
    # 无论 status：due_at 是唯一时间门闩
    if due > current:
        return due
    # 已到期：未完成要续跑；done 也表示下一轮已到期
    return due


def resume_stat_date(repo: Repository, name: str, *, fallback: str) -> str:
    """续跑时优先用状态里的归属日，避免跨日重启跑错天。"""
    state = load_state(repo, name)
    if state is not None and state.stat_date and state.status in _ACTIVE:
        return state.stat_date
    return fallback