from __future__ import annotations

import os
import signal
import subprocess
import threading
import uuid
from pathlib import Path

from . import db
from .commands import CommandSpec
from .config import settings
from .modules import build_module_task
from .schemas import JobCreate

_running: dict[str, subprocess.Popen] = {}
_lock = threading.Lock()


def _build_task(request: JobCreate) -> CommandSpec:
    conda_env = request.conda_env or settings.task_conda_env
    spec = build_module_task(request.type, request.params, conda_env)
    if spec is None:
        raise ValueError(f"unsupported job type: {request.type}")
    env = dict(spec.env)
    if request.gpus:
        env["CUDA_VISIBLE_DEVICES"] = ",".join(str(gpu) for gpu in request.gpus)
    return CommandSpec(
        command=spec.command,
        cwd=spec.cwd,
        description=spec.description,
        env=env,
        stdin_text=spec.stdin_text,
    )


def create_job(request: JobCreate) -> dict:
    spec = _build_task(request)
    job_id = uuid.uuid4().hex
    settings.log_dir.mkdir(parents=True, exist_ok=True)
    log_path = settings.log_dir / f"{job_id}.log"
    job = {
        "id": job_id,
        "type": request.type,
        "status": "queued",
        "params": request.params,
        "command": spec.command,
        "cwd": str(spec.cwd),
        "description": spec.description,
        "pid": None,
        "exit_code": None,
        "created_at": db.utc_now(),
        "started_at": None,
        "finished_at": None,
        "log_path": str(log_path),
        "error": None,
    }
    db.insert_job(job)
    thread = threading.Thread(target=_run_job, args=(job_id, spec, log_path), daemon=True)
    thread.start()
    return db.get_job(job_id)


def _run_job(job_id: str, spec: CommandSpec, log_path: Path) -> None:
    env = os.environ.copy()
    env.update(spec.env)
    db.update_job(job_id, status="running", started_at=db.utc_now())
    try:
        with log_path.open("ab") as log_file:
            log_file.write(("COMMAND: " + " ".join(spec.command) + "\n").encode("utf-8"))
            log_file.flush()
            process = subprocess.Popen(
                spec.command,
                cwd=str(spec.cwd),
                env=env,
                stdin=subprocess.PIPE if spec.stdin_text is not None else None,
                stdout=log_file,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
            with _lock:
                _running[job_id] = process
            db.update_job(job_id, pid=process.pid)
            if spec.stdin_text is not None and process.stdin is not None:
                process.stdin.write(spec.stdin_text.encode("utf-8"))
                process.stdin.close()
            exit_code = process.wait()
        with _lock:
            _running.pop(job_id, None)
        current = db.get_job(job_id)
        if current and current["status"] == "cancelled":
            db.update_job(job_id, exit_code=exit_code, finished_at=db.utc_now())
        elif exit_code == 0:
            db.update_job(job_id, status="success", exit_code=exit_code, finished_at=db.utc_now())
        else:
            db.update_job(job_id, status="failed", exit_code=exit_code, finished_at=db.utc_now())
    except Exception as exc:
        with _lock:
            _running.pop(job_id, None)
        db.update_job(job_id, status="failed", error=str(exc), finished_at=db.utc_now())


def cancel_job(job_id: str) -> dict | None:
    job = db.get_job(job_id)
    if not job:
        return None
    db.update_job(job_id, status="cancelled", finished_at=db.utc_now())
    with _lock:
        process = _running.get(job_id)
    if process and process.poll() is None:
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
    return db.get_job(job_id)