"""One worker thread per GPU, pulling jobs off the shared queue."""

import threading
import time
import traceback

from . import db
from .trainer import run_training

_cancel = set()
_cancel_lock = threading.Lock()
_workers = []


def request_cancel(job_id):
    with _cancel_lock:
        _cancel.add(job_id)
    j = db.get_job(job_id)
    if j and j["status"] == "queued":
        db.update_job(job_id, status="cancelled", finished_at=time.time())
        return True
    return bool(j)


def _is_cancelled(job_id):
    with _cancel_lock:
        return job_id in _cancel


_devices_cache = None


def devices():
    """Never raises and never blocks twice — CUDA probing is done once."""
    global _devices_cache
    if _devices_cache is not None:
        return _devices_cache
    found = ["cpu"]
    try:
        import torch
        if torch.cuda.is_available() and torch.cuda.device_count() > 0:
            found = [f"cuda:{i}" for i in range(torch.cuda.device_count())]
    except Exception as exc:
        print(f"[worker] no CUDA ({exc!r}) — falling back to CPU", flush=True)
    _devices_cache = found
    return found


def _loop(gpu_index, device, tokenizer_dir):
    while True:
        try:
            job = db.claim_next_queued(gpu_index)
        except Exception:
            job = None
        if not job:
            time.sleep(2)
            continue

        jid = job["id"]

        def log(msg, _jid=jid, _dev=device):
            print(f"[{_dev}][{_jid}] {msg}", flush=True)
            try:
                db.add_log(_jid, msg)
            except Exception:
                pass

        # Don't start on a barely-filled cache, or a 300M-token run would
        # loop over the same few million tokens hundreds of times.
        try:
            from . import data as ndata
            need = min(int(job["target_tokens"]), ndata.TARGET_TOKENS)
            last_note = 0.0
            while ndata.available_tokens() < need and ndata.state()["status"] == "building":
                if time.time() - last_note > 60:
                    log(f"waiting for the fineweb-edu cache to reach "
                        f"{need / 1e6:.0f}M tokens "
                        f"({ndata.available_tokens() / 1e6:.0f}M ready so far — "
                        f"this is a one-time cost per container)")
                    last_note = time.time()
                if _is_cancelled(jid):
                    break
                time.sleep(5)
        except Exception:
            pass

        if _is_cancelled(jid):
            db.update_job(jid, status="cancelled", finished_at=time.time())
            with _cancel_lock:
                _cancel.discard(jid)
            continue

        log(f"starting on {device}")
        try:
            run_training(job, device, tokenizer_dir, lambda: _is_cancelled(jid), log)
        except Exception as exc:
            tb = traceback.format_exc(limit=6)
            print(tb, flush=True)
            log(f"FAILED: {exc!r}")
            db.update_job(jid, status="failed", error=str(exc)[:500],
                          finished_at=time.time())
        finally:
            with _cancel_lock:
                _cancel.discard(jid)
            try:
                import torch
                if device.startswith("cuda"):
                    torch.cuda.empty_cache()
            except Exception:
                pass


def start_workers(tokenizer_dir):
    global _workers
    if _workers:
        return _workers
    for i, dev in enumerate(devices()):
        if dev.startswith("cuda"):
            try:
                import torch
                name = torch.cuda.get_device_name(i)
                gb = torch.cuda.get_device_properties(i).total_memory / 1024**3
                print(f"[worker] {dev}: {name} ({gb:.0f} GB)", flush=True)
            except Exception:
                pass
        t = threading.Thread(target=_loop, args=(i, dev, tokenizer_dir),
                             name=f"worker-{dev}", daemon=True)
        t.start()
        _workers.append(t)
        print(f"[worker] started on {dev}", flush=True)
    return _workers


def worker_info():
    devs = devices()
    running = {j["gpu_index"]: j for j in db.list_jobs(status=["running"], limit=64)}
    return [{"index": i, "device": d, "job": running.get(i)} for i, d in enumerate(devs)]