"""Per-group task queues and subprocess execution."""

from __future__ import annotations

import json
import os
import signal
import shutil
import subprocess
import threading
import time
from collections.abc import Mapping
from pathlib import Path

from .config import Settings
from .discovery import BackgroundTask, Group, Task, TaskCatalog
from .inputs import PendingUpload
from .store import Run, RunStore


class RunManager:
    def __init__(self, settings: Settings, catalog: TaskCatalog, store: RunStore) -> None:
        self.settings = settings
        self.catalog = catalog
        self.store = store
        self._wake_events: dict[str, threading.Event] = {}
        self._workers: dict[str, threading.Thread] = {}
        self._processes: dict[int, subprocess.Popen[bytes]] = {}
        self._lock = threading.Lock()

    def run_task(
        self,
        task: Task,
        input_data: dict[str, str | int | list[str]] | None = None,
        uploads: Mapping[str, PendingUpload] | None = None,
    ) -> tuple[Run, bool]:
        run, created = self.store.enqueue(
            target_kind="task",
            target_path=task.id,
            queue_group=task.group_id,
            input_data=input_data,
        )
        if created:
            try:
                uploaded_inputs = self._save_uploads(run.id, uploads or {})
            except (OSError, ValueError) as error:
                self._delete_uploads(run.id)
                return self.store.finish(
                    run.id,
                    status="failed",
                    message=f"Could not save uploaded file: {error}",
                ), True
            if uploaded_inputs:
                run = self.store.set_input(run.id, {**(input_data or {}), **uploaded_inputs})
            self._wake_worker(task.group_id)
        return run, created

    def run_group(self, group: Group) -> tuple[Run, bool]:
        run, created = self.store.enqueue(
            target_kind="group",
            target_path=group.id,
            queue_group=group.id,
        )
        if created:
            self._wake_worker(group.id)
        return run, created

    def request_cancel(self, run_id: int) -> Run | None:
        run = self.store.request_cancel(run_id)
        if run is None:
            return None
        if run.parent_run_id is None and run.target_kind == "group":
            child = self.store.active_child(run.id)
            if child is not None:
                self.request_cancel(child.id)
        self._terminate_active_process(run.id)
        completed = self.store.get(run.id)
        if completed is not None and completed.status == "cancelled":
            self._delete_uploads(run.id)
        return completed

    def artifacts(self, run_id: int) -> list[str]:
        root = self._artifact_root(run_id)
        if not root.is_dir():
            return []
        return [
            path.name
            for path in sorted(root.iterdir(), key=lambda item: item.name)
            if path.is_file() and not path.is_symlink()
        ]

    def artifact_path(self, run_id: int, filename: str) -> Path | None:
        path = self._artifact_root(run_id) / filename
        if Path(filename).name != filename or path.is_symlink() or not path.is_file():
            return None
        return path

    def delete_artifact(self, run_id: int, filename: str) -> None:
        path = self.artifact_path(run_id, filename)
        if path is None:
            return
        path.unlink(missing_ok=True)
        try:
            path.parent.rmdir()
        except OSError:
            pass

    def cleanup_transient_files(self, *, artifact_max_age_seconds: int = 86_400) -> None:
        """Remove uploads from interrupted runs and expired undownloaded artifacts."""

        runs_root = self.settings.state_root / "runs"
        if not runs_root.is_dir():
            return
        cutoff = time.time() - artifact_max_age_seconds
        for path in runs_root.iterdir():
            if not path.is_dir():
                continue
            run_id_text, separator, kind = path.name.partition(".")
            if not separator or not run_id_text.isdigit() or kind not in {"uploads", "artifacts"}:
                continue
            run = self.store.get(int(run_id_text))
            if kind == "uploads" and (run is None or run.status not in {"queued", "running"}):
                shutil.rmtree(path, ignore_errors=True)
            elif kind == "artifacts":
                try:
                    expired = path.stat().st_mtime < cutoff
                except FileNotFoundError:
                    continue
                if expired:
                    shutil.rmtree(path, ignore_errors=True)

    def run_background_task(self, task: BackgroundTask) -> None:
        """Run a declared internal task without creating high-volume run history."""

        log_path = self.settings.state_root / "background" / f"{task.id}.log"
        log_path.parent.mkdir(parents=True, exist_ok=True)
        environment = self._environment(task.group_id)
        try:
            with log_path.open("wb") as log_file:
                process = subprocess.Popen(
                    [self.settings.python_executable, str(task.path)],
                    cwd=task.path.parent,
                    stdin=subprocess.DEVNULL,
                    stdout=log_file,
                    stderr=subprocess.STDOUT,
                    env=environment,
                    start_new_session=True,
                )
                exit_code = process.wait()
                if exit_code:
                    log_file.write(f"Background task exited with status {exit_code}.\n".encode())
        except OSError as error:
            log_path.write_text(f"Could not start background task: {error}\n", encoding="utf-8")

    def _wake_worker(self, group_id: str) -> None:
        with self._lock:
            event = self._wake_events.setdefault(group_id, threading.Event())
            worker = self._workers.get(group_id)
            if worker is None or not worker.is_alive():
                worker = threading.Thread(
                    target=self._worker_loop,
                    args=(group_id, event),
                    name=f"task-group-{group_id}",
                    daemon=True,
                )
                self._workers[group_id] = worker
                worker.start()
            event.set()

    def _worker_loop(self, group_id: str, event: threading.Event) -> None:
        while True:
            event.clear()
            run = self.store.claim_next(group_id)
            if run is None:
                event.wait()
                continue
            self._execute_root_run(run)

    def _execute_root_run(self, run: Run) -> None:
        if run.target_kind == "task":
            task = self.catalog.task(run.target_path)
            if task is None:
                self.store.finish(
                    run.id,
                    status="failed",
                    message="Task is no longer available.",
                )
                return
            self._execute_task(run, task)
            return

        group = self.catalog.group(run.target_path)
        if group is None:
            self.store.finish(
                run.id,
                status="failed",
                message="Group is no longer available.",
            )
            return
        self._execute_group(run, group)

    def _execute_group(self, parent_run: Run, group: Group) -> None:
        for task in group.tasks:
            if self.store.is_cancel_requested(parent_run.id):
                self.store.finish(
                    parent_run.id,
                    status="cancelled",
                    message="Group cancellation requested.",
                )
                return
            child = self.store.create_child(parent_run, task.id)
            completed_child = self._execute_task(child, task, parent_run_id=parent_run.id)
            if completed_child.status != "succeeded":
                parent_status = (
                    "cancelled"
                    if completed_child.status == "cancelled"
                    or self.store.is_cancel_requested(parent_run.id)
                    else "failed"
                )
                self.store.finish(
                    parent_run.id,
                    status=parent_status,
                    message=f"Stopped after {task.id}: {completed_child.status}.",
                )
                return
        self.store.finish(parent_run.id, status="succeeded")

    def _execute_task(
        self,
        run: Run,
        task: Task,
        *,
        parent_run_id: int | None = None,
    ) -> Run:
        if self.store.is_cancel_requested(run.id) or (
            parent_run_id is not None and self.store.is_cancel_requested(parent_run_id)
        ):
            return self.store.finish(run.id, status="cancelled", message="Cancelled before execution.")

        log_path = self.settings.state_root / "runs" / f"{run.id}.log"
        log_path.parent.mkdir(parents=True, exist_ok=True)
        input_path = self.settings.state_root / "runs" / f"{run.id}.input.json"
        artifact_root = self._artifact_root(run.id)
        artifact_root.mkdir(parents=True, exist_ok=True)
        try:
            input_path.write_text(run.input_json + "\n", encoding="utf-8")
            input_path.chmod(0o600)
        except OSError as error:
            self._delete_uploads(run.id)
            shutil.rmtree(artifact_root, ignore_errors=True)
            return self.store.finish(
                run.id,
                status="failed",
                message=f"Could not prepare task input: {error}",
            )
        environment = self._environment(task.group_id)
        environment["SERVER_MAINTENANCE_INPUT"] = str(input_path)
        environment["SERVER_MAINTENANCE_ARTIFACTS"] = str(artifact_root)

        try:
            with log_path.open("wb") as log_file:
                process = subprocess.Popen(
                    [self.settings.python_executable, str(task.path)],
                    cwd=task.path.parent,
                    stdin=subprocess.DEVNULL,
                    stdout=log_file,
                    stderr=subprocess.STDOUT,
                    env=environment,
                    start_new_session=True,
                )
                with self._lock:
                    self._processes[run.id] = process
                exit_code = process.wait()
        except OSError as error:
            self._delete_uploads(run.id)
            shutil.rmtree(artifact_root, ignore_errors=True)
            return self.store.finish(
                run.id,
                status="failed",
                message=f"Could not start task: {error}",
                log_path=str(log_path),
            )
        finally:
            with self._lock:
                self._processes.pop(run.id, None)
            self._delete_uploads(run.id)

        if self.store.is_cancel_requested(run.id) or (
            parent_run_id is not None and self.store.is_cancel_requested(parent_run_id)
        ):
            shutil.rmtree(artifact_root, ignore_errors=True)
            return self.store.finish(
                run.id,
                status="cancelled",
                exit_code=exit_code,
                log_path=str(log_path),
            )
        completed = self.store.finish(
            run.id,
            status="succeeded" if exit_code == 0 else "failed",
            exit_code=exit_code,
            log_path=str(log_path),
        )
        if completed.status != "succeeded":
            shutil.rmtree(artifact_root, ignore_errors=True)
        else:
            try:
                artifact_root.rmdir()
            except OSError:
                pass
        self.cleanup_transient_files()
        return completed

    def _environment(self, group_id: str) -> dict[str, str]:
        task_state = self.settings.state_root / "tasks" / group_id
        task_state.mkdir(parents=True, exist_ok=True)
        environment = {
            "PATH": "/usr/local/bin:/usr/bin:/bin",
            "PYTHONUNBUFFERED": "1",
            "SERVER_MAINTENANCE_CREDENTIALS": str(self.settings.credentials_root),
            "SERVER_MAINTENANCE_STATE": str(task_state),
        }
        if self.settings.timezone_name:
            environment["TZ"] = self.settings.timezone_name
        return environment

    def _upload_root(self, run_id: int) -> Path:
        return self.settings.state_root / "runs" / f"{run_id}.uploads"

    def _artifact_root(self, run_id: int) -> Path:
        return self.settings.state_root / "runs" / f"{run_id}.artifacts"

    def _save_uploads(
        self,
        run_id: int,
        uploads: Mapping[str, PendingUpload],
    ) -> dict[str, str]:
        if not uploads:
            return {}
        root = self._upload_root(run_id)
        root.mkdir(parents=True, exist_ok=False)
        saved = {}
        try:
            for field_name, upload in uploads.items():
                suffix = upload.extension or ".upload"
                destination = root / f"{field_name}{suffix}"
                temporary = destination.with_suffix(destination.suffix + ".part")
                total = 0
                with temporary.open("xb") as output:
                    while chunk := upload.stream.read(64 * 1024):
                        total += len(chunk)
                        if upload.maximum_bytes is not None and total > upload.maximum_bytes:
                            raise ValueError(f"{field_name} exceeds its size limit")
                        output.write(chunk)
                temporary.replace(destination)
                destination.chmod(0o600)
                saved[field_name] = str(destination)
        except Exception:
            shutil.rmtree(root, ignore_errors=True)
            raise
        return saved

    def _delete_uploads(self, run_id: int) -> None:
        shutil.rmtree(self._upload_root(run_id), ignore_errors=True)

    def _terminate_active_process(self, run_id: int) -> None:
        with self._lock:
            process = self._processes.get(run_id)
        if process is None or process.poll() is not None:
            return
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            return
        threading.Thread(
            target=self._kill_after_grace,
            args=(process,),
            name=f"task-cancel-{run_id}",
            daemon=True,
        ).start()

    def _kill_after_grace(self, process: subprocess.Popen[bytes]) -> None:
        try:
            process.wait(timeout=self.settings.cancel_grace_seconds)
        except subprocess.TimeoutExpired:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass