"""Huey task definitions and shared batch state for the admin panel."""
from __future__ import annotations

import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any

from huey import MemoryHuey

huey = MemoryHuey("seed", immediate=False)

# ── Shared state (all accessed under _lock) ───────────────────────────────────
_lock = threading.Lock()

_current_batch_id: str | None = None
_current_result: Any = None          # Huey Result — kept for revocation

_batch_meta: dict[str, Any] = {}     # batch_id → aggregate metadata
_worker_slots: dict[str, dict] = {}  # f"{batch_id}:{slot:04d}" → worker info
_stop_events: dict[str, threading.Event] = {}


# ── Public helpers ─────────────────────────────────────────────────────────────

def get_current_batch_id() -> str | None:
    with _lock:
        return _current_batch_id


def set_current(batch_id: str, result: Any) -> None:
    global _current_batch_id, _current_result
    with _lock:
        _current_batch_id = batch_id
        _current_result = result


def stop_current() -> bool:
    """Signal the running batch to stop. Returns True if there was something to stop."""
    with _lock:
        bid = _current_batch_id
        result = _current_result

    if bid is None:
        return False

    # Revoke task if it hasn't started yet (still pending in Huey queue)
    if result is not None:
        try:
            result.revoke(revoke_once=True)
        except Exception:
            pass

    # Signal running workers to stop after their current step
    ev = _stop_events.get(bid)
    if ev:
        ev.set()

    with _lock:
        if bid in _batch_meta:
            _batch_meta[bid]["stopping"] = True

    return True


def batch_status() -> dict[str, Any]:
    with _lock:
        bid = _current_batch_id
        if bid is None:
            return {"batch_id": None, "running": False}

        meta = dict(_batch_meta.get(bid, {}))
        prefix = f"{bid}:"
        workers = [dict(v) for k, v in _worker_slots.items() if k.startswith(prefix)]

    meta["batch_id"] = bid
    meta["workers_detail"] = sorted(workers, key=lambda w: w.get("slot", 0))
    meta["active_workers"] = sum(1 for w in workers if w.get("status") == "running")
    return meta


# ── Huey task ─────────────────────────────────────────────────────────────────

@huey.task(name="run_batch")
def run_batch_task(
    batch_id: str,
    cfg: dict[str, Any],
    workers: int,
    total_runs: int,
    headless: bool,
    stagger_ms: int,
) -> dict[str, Any]:
    import asyncio

    from crawler.driver import driver_session
    from crawler.dynamic_flow import DynamicFlow
    from crawler.scenarios.fresh_session import FreshSessionScenario
    from db.repositories.runs import runs_repo
    from logger import get_logger

    log = get_logger(__name__)

    stop_ev = threading.Event()
    with _lock:
        _stop_events[batch_id] = stop_ev
        _batch_meta[batch_id] = {
            "running": True,
            "stopping": False,
            "total": total_runs,
            "completed": 0,
            "succeeded": 0,
            "failed": 0,
            "configured_workers": workers,
            "started_at": time.time(),
            "finished_at": None,
            "error": None,
        }

    def _run_one(slot: int) -> bool:
        slot_key = f"{batch_id}:{slot:04d}"

        if stop_ev.is_set():
            with _lock:
                _worker_slots[slot_key] = {
                    "slot": slot, "status": "skipped",
                    "started_at": None, "finished_at": None, "duration_ms": None, "error": "stopped",
                }
            return False

        with _lock:
            _worker_slots[slot_key] = {
                "slot": slot, "status": "running",
                "started_at": time.time(), "finished_at": None, "duration_ms": None, "error": None,
            }

        t0 = time.monotonic()
        identifier = f"batch-{batch_id[:8]}-{slot:04d}"
        log.info("[batch:%s] slot %d starting", batch_id[:8], slot)

        try:
            target_url: str = cfg.get("target_url") or ""
            with driver_session(headless=headless) as driver:
                FreshSessionScenario(driver, target_url).run()
                result = DynamicFlow(driver, cfg).run() if cfg else None

            total_ms = int((time.monotonic() - t0) * 1000)
            success = result.success if result else True
            reason = result.failure_reason if result else None
            steps = [
                {"name": s.name, "success": s.success, "error": s.error, "duration_ms": s.duration_ms}
                for s in (result.steps if result else [])
            ]
            asyncio.run(runs_repo.insert("batch", identifier, success, reason, steps, total_ms))

            status = "ok" if success else "failed"
            with _lock:
                _worker_slots[slot_key].update({
                    "status": status, "finished_at": time.time(), "duration_ms": total_ms,
                })
            log.info("[batch:%s] slot %d %s — %d ms", batch_id[:8], slot, status, total_ms)
            return success

        except Exception as exc:
            total_ms = int((time.monotonic() - t0) * 1000)
            asyncio.run(runs_repo.insert("batch", identifier, False, str(exc), [], total_ms))
            with _lock:
                _worker_slots[slot_key].update({
                    "status": "error", "error": str(exc),
                    "finished_at": time.time(), "duration_ms": total_ms,
                })
            log.error("[batch:%s] slot %d error — %s", batch_id[:8], slot, exc)
            return False

    succeeded = 0
    failed = 0

    try:
        with ThreadPoolExecutor(
            max_workers=workers,
            thread_name_prefix=f"batch-{batch_id[:8]}",
        ) as pool:
            futures = []
            for i in range(total_runs):
                if stop_ev.is_set():
                    break
                futures.append(pool.submit(_run_one, i))
                if stagger_ms > 0 and i < total_runs - 1:
                    if stop_ev.wait(timeout=stagger_ms / 1000):
                        break  # stop was signalled during the stagger delay

            for fut in as_completed(futures):
                try:
                    ok = fut.result()
                except Exception:
                    ok = False
                if ok:
                    succeeded += 1
                else:
                    failed += 1
                with _lock:
                    _batch_meta[batch_id]["completed"] = succeeded + failed
                    _batch_meta[batch_id]["succeeded"] = succeeded
                    _batch_meta[batch_id]["failed"] = failed

    except Exception as exc:
        with _lock:
            _batch_meta[batch_id]["error"] = str(exc)

    finally:
        with _lock:
            _batch_meta[batch_id]["running"] = False
            _batch_meta[batch_id]["finished_at"] = time.time()

    return {"batch_id": batch_id, "succeeded": succeeded, "failed": failed}