"""Job execution in separate processes (spawned, so each run has its own interpreter,
thread caps and model state). Progress is written to the SQLite job table, throttled."""
from __future__ import annotations

import json
import os
import shutil
import time
import traceback
import zipfile
from concurrent.futures import ProcessPoolExecutor
from multiprocessing import get_context
from pathlib import Path

from . import db


class Cancelled(Exception):
    pass


def _bundle(out: Path, dest: Path):
    tmp = dest.with_suffix(".tmp")
    with zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED) as z:
        for f in sorted(out.rglob("*")):
            if f.is_file():
                z.write(f, Path("ppi_benchmark_evaluation") / f.relative_to(out))
    tmp.replace(dest)


def run_job(job_id: str, db_path: str, dataset_dir: str, job_dir: str, config: dict, threads: int):
    """Entry point inside the worker process."""
    for var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
        os.environ[var] = str(threads)
    db_path, dataset_dir, job_dir = Path(db_path), Path(dataset_dir), Path(job_dir)
    from ppi_controls import Inputs, RunConfig, run_pipeline
    from ppi_controls.io import InputError

    meta = json.loads((dataset_dir / "meta.json").read_text())
    files = {k: dataset_dir / v["stored"] for k, v in meta["files"].items()}
    inputs = Inputs(files["proteins"], files["pairs"], files.get("species"), files.get("clusters"), files.get("embeddings"),
                    mappings=meta.get("mappings", {}), labels={k: v["name"] for k, v in meta["files"].items()})
    cfg = RunConfig(preset=config.get("preset", "quick"), seed=int(config.get("seed", 42)),
                    species_mode=config.get("species_mode", "auto"), n_jobs=threads,
                    embeddings_name=config.get("embeddings_name"))
    db.update(db_path, job_id, status="running", started=time.time(), message="Starting")
    last = {"t": 0.0}

    def progress(ev):
        now = time.time()
        if now - last["t"] < 0.4 and ev["stage_fraction"] not in (0.0, 1.0):
            return
        last["t"] = now
        job = db.get(db_path, job_id)
        if job and job["cancel"]:
            raise Cancelled()
        db.update(db_path, job_id, stage=ev["stage"], stage_label=ev["stage_label"], progress=ev["progress"],
                  message=ev["message"])

    out = job_dir / "output"
    try:
        shutil.rmtree(out, ignore_errors=True)
        run_pipeline(inputs, cfg, out, progress)
        db.update(db_path, job_id, message="Packaging results", progress=0.995)
        _bundle(out, job_dir / "bundle.zip")
        db.update(db_path, job_id, status="done", progress=1.0, message="Done", finished=time.time())
    except Cancelled:
        db.update(db_path, job_id, status="cancelled", message="Cancelled", finished=time.time())
    except InputError as e:
        db.update(db_path, job_id, status="failed", finished=time.time(), message="Input problem",
                  error={"message": str(e), "issues": [i.to_dict() for i in e.issues]})
    except Exception as e:  # noqa: BLE001 - reported to the user, logged in full
        (job_dir / "error.log").write_text(traceback.format_exc())
        db.update(db_path, job_id, status="failed", finished=time.time(), message="Analysis failed",
                  error={"message": f"{type(e).__name__}: {e}"})


class JobManager:
    def __init__(self, settings):
        self.settings = settings
        self.pool = ProcessPoolExecutor(max_workers=settings.max_concurrent_jobs, mp_context=get_context("spawn"))

    def submit(self, job_id: str, dataset_dir: Path, job_dir: Path, config: dict):
        fut = self.pool.submit(run_job, job_id, str(self.settings.db_path), str(dataset_dir), str(job_dir),
                               config, self.settings.job_threads)

        def done(f):
            exc = f.exception()
            if exc is not None:  # the worker process itself died (e.g. out of memory)
                job = db.get(self.settings.db_path, job_id)
                if job and job["status"] not in db.TERMINAL:
                    db.update(self.settings.db_path, job_id, status="failed", finished=time.time(),
                              error={"message": f"Worker crashed: {type(exc).__name__}: {exc}"})
        fut.add_done_callback(done)

    def shutdown(self):
        self.pool.shutdown(wait=False, cancel_futures=True)