"""Persistent batch coordinator layered on the existing device queue."""

import asyncio
import logging
import os
import time

from apps.admin_console.core.state import state
from apps.admin_console.services.automation_store import store, new_id
from apps.admin_console.services.task_import import case_goal

logger = logging.getLogger(__name__)
TERMINAL = {"completed", "failed", "cancelled", "skipped", "interrupted"}
background_tasks: set[asyncio.Task] = set()


def spawn(coro):
    task = asyncio.create_task(coro)
    background_tasks.add(task)
    task.add_done_callback(background_tasks.discard)
    return task


def register_submission(item: dict, case: dict | None = None, formats=None, **extra):
    sid = str(item["session_id"])
    if store.job(sid):
        return
    store.save_job(
        {
            "session_id": sid,
            "case": case or {"name": item.get("goal", sid)[:120], "goal": item.get("goal", sid)},
            "status": "queued",
            "created_at": time.time(),
            "device_serial": item.get("device_serial"),
            "formats": formats if formats is not None else store.preferences()["formats"],
            **extra,
        }
    )


def recover_batches():
    """Recovery never silently repeats a potentially side-effecting interrupted task."""
    for job in store.jobs():
        if job["status"] in {"running", "queued", "finalizing"}:
            job.update(
                status=job.get("terminal_status", "interrupted"),
                error=job.get("error")
                or ("" if job.get("terminal_status") else "服务中断；未完成任务需手动重跑"),
                ended_at=job.get("ended_at") or time.time(),
            )
            store.save_job(job)
    for batch in store.records("batch"):
        if has_legacy_replay(batch["id"]):
            for job in store.jobs(batch["id"]):
                if job["status"] not in TERMINAL:
                    job.update(status="cancelled", error="操作回放功能已移除", ended_at=time.time())
                    store.save_job(job)
            if batch["status"] not in TERMINAL:
                batch.update(status="cancelled", reason="操作回放功能已移除；保留历史供查看")
                store.put("batch", batch["id"], batch)
            continue
        if batch["status"] in {"running", "pausing", "paused"}:
            batch.update(
                status="paused", reason="服务已重启；请检查中断任务，点击继续执行未开始任务"
            )
            for job in store.jobs(batch["id"]):
                if job["status"] in {"running", "queued", "finalizing"}:
                    job.update(
                        status="interrupted", error="服务中断，需手动重跑", ended_at=time.time()
                    )
                    store.save_job(job)
            store.put("batch", batch["id"], batch)


async def on_terminal(item: dict, status: str):
    from apps.admin_console.services.report_export import export_session, session_report

    sid = str(item["session_id"])
    job = store.job(sid)
    if not job:
        register_submission(item)
        job = store.job(sid)
    if job is None:
        raise RuntimeError("任务元数据写入失败")
    job.update(status="finalizing", ended_at=time.time())
    if job.get("timed_out"):
        status = "failed"
    store.save_job(job)
    try:
        job["status"] = status
        report = await asyncio.to_thread(session_report, sid)
        if report["verdict"] == "验证失败" or report["outcome"].get("task_status") in {
            "blocked",
            "partial",
            "failed",
        }:
            job["status"] = "failed"
            job["error"] = "任务执行结束，但业务检查未通过"
        # Persist terminal classification before export, but the coordinator waits for finalization.
        store.save_job({**job, "status": "finalizing", "terminal_status": job["status"]})
        await asyncio.to_thread(export_session, sid, job["formats"])
        job["report_ready"] = True
    except Exception as exc:
        logger.exception("Report finalization failed for %s", sid)
        job["report_error"] = str(exc)[:500]
    store.save_job(job)


async def tick_batches():
    from apps.admin_console.server import task_queue_service
    from artemis.runtime import device_pool
    from apps.admin_console.services.report_export import export_batch

    for batch in store.records("batch"):
        if batch["status"] not in {"running", "paused"}:
            continue
        jobs = store.jobs(batch["id"])
        active = [j for j in jobs if j["status"] in {"running", "queued", "finalizing"}]
        if active:
            if state.is_paused:
                batch.update(status="paused", reason="模型请求已暂停，请先恢复模型连接再继续批次")
                store.put("batch", batch["id"], batch)
            for job in active:
                limit = int(batch.get("timeout_seconds", 1800))
                if (
                    job["status"] != "finalizing"
                    and not job.get("timed_out")
                    and time.time() - job.get("started_at", time.time()) > limit
                ):
                    job.update(error="任务达到批次配置的执行时限", timed_out=True)
                    store.save_job(job)
                    task_queue_service.stop_tasks(session_id=job["session_id"])
                    batch.update(status="paused", reason="任务超时，已暂停剩余任务；请确认设备状态")
                    store.put("batch", batch["id"], batch)
            continue
        if batch["status"] != "running":
            continue
        pending = [j for j in jobs if j["status"] == "pending"]
        failed = [j for j in jobs if j["status"] in {"failed", "cancelled", "interrupted"}]
        if batch.get("stop_on_failure") and failed and pending:
            for job in pending:
                job.update(status="skipped", error="本批次设置为失败即停止", ended_at=time.time())
                store.save_job(job)
            pending = []
        if not pending:
            batch.update(status="completed", ended_at=time.time())
            try:
                await asyncio.to_thread(export_batch, batch["id"])
                batch["report_ready"] = True
            except Exception as exc:
                logger.exception("Batch report generation failed")
                batch["report_error"] = str(exc)[:500]
            store.put("batch", batch["id"], batch)
            continue
        if state.is_paused:
            batch.update(status="paused", reason="模型请求已暂停，请先恢复模型连接再继续批次")
            store.put("batch", batch["id"], batch)
            continue
        serial = batch.get("device_serial")
        devices = await device_pool.try_list_devices_async()
        rejection = (
            None
            if serial and any(d.serial == serial and d.state == "device" for d in devices or [])
            else "无法确认所选设备在线且已授权，请恢复设备连接后继续批次"
        )
        if rejection:
            batch.update(status="paused", reason=str(rejection))
            store.put("batch", batch["id"], batch)
            continue
        if (store.get("batch", batch["id"]) or {}).get("status") != "running":
            continue
        job = pending[0]
        job.update(status="queued", started_at=time.time())
        store.save_job(job)
        try:
            response = await task_queue_service.enqueue_tasks(
                [case_goal(job["case"])],
                profile=batch.get("profile", "flash"),
                expected_output="用简体中文报告实际结果与证据。明确区分执行完成、验证通过、验证失败和前置条件未满足。"
                + job["case"].get("expected", ""),
                enable_outputter=True,
                device_serial=serial,
                session_id=job["session_id"],
                ingress="batch",
            )
            if response.get("status") == "rejected" or not response.get("tasks"):
                raise ValueError(response.get("error") or "队列拒绝接收任务")
            bound = response["tasks"][0].get("device_serial")
            if bound and not serial:
                batch["device_serial"] = bound
                store.put("batch", batch["id"], batch)
        except Exception as exc:
            job.update(status="pending", error=str(exc)[:500])
            store.save_job(job)
            batch.update(status="paused", reason=str(exc)[:500])
            store.put("batch", batch["id"], batch)


async def batch_loop():
    recover_batches()
    while True:
        try:
            await tick_batches()
        except Exception:
            logger.exception("Batch coordinator failed; persisted work is retained")
        await asyncio.sleep(1)


def has_legacy_replay(batch_id: str) -> bool:
    return any(
        j.get("workflow_id") or j.get("execution_mode") == "strict" or j.get("adaptive_source")
        for j in store.jobs(batch_id)
    )


def create_batch(payload: dict) -> dict:
    bid = new_id()
    batch = {
        "id": bid,
        "name": payload["name"],
        "status": "draft",
        "created_at": time.time(),
        "formats": payload.get("formats") or store.preferences()["formats"],
        "profile": payload.get("profile", "flash"),
        "device_serial": payload.get("device_serial"),
        "timeout_seconds": payload.get("timeout_seconds")
        or int(os.getenv("ARTEMIS_BATCH_TASK_TIMEOUT", "1800")),
        "stop_on_failure": payload.get("stop_on_failure", False),
        "request_id": payload.get("request_id"),
    }
    jobs = []
    for number, case in enumerate(payload["cases"], 1):
        jobs.append(
            {
                "session_id": new_id(),
                "batch_id": bid,
                "position": number,
                "case": case,
                "status": "pending",
                "formats": batch["formats"],
                "created_at": time.time(),
                "device_serial": batch["device_serial"],
            }
        )
    return store.create_batch(batch, jobs, payload)


def batch_action(batch_id: str, action: str) -> dict:
    batch = store.get("batch", batch_id)
    if not batch:
        raise ValueError("批次不存在")
    if has_legacy_replay(batch_id):
        raise ValueError("操作回放功能已移除；历史记录仅供查看")
    if action == "pause":
        if batch["status"] not in {"running", "paused"}:
            raise ValueError("只能暂停正在执行的批次")
        batch.update(status="paused", reason="用户暂停；当前任务继续，后续任务等待")
    elif action == "resume":
        if batch["status"] not in {"draft", "paused"}:
            raise ValueError("只能开始草稿或继续已暂停批次")
        batch.update(status="running", reason="")
    elif action == "retry_failed":
        if batch["status"] == "running":
            raise ValueError("请先暂停批次")
        failed = [
            j for j in store.jobs(batch_id) if j["status"] in {"failed", "cancelled", "interrupted"}
        ]
        if not failed:
            raise ValueError("没有可重跑的失败或中断任务")
        revised = create_batch(
            {**batch, "name": batch["name"] + " · 失败项重跑", "cases": [j["case"] for j in failed]}
        )
        return revised
    else:
        raise ValueError("不支持的批次操作")
    return store.put("batch", batch_id, batch)