"""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


ORPHANED_TASK_STATE_FILENAME = ".orphaned-groups.json"
ORPHANED_TASK_STATE_RETENTION_SECONDS = 14 * 86400


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]:
        sensitive_names = {field.name for field in task.inputs if field.sensitive}
        stored_input = self._redact_input(input_data or {}, sensitive_names)
        run, created = self.store.enqueue(
            target_kind="task",
            target_path=task.id,
            queue_group=task.group_id,
            input_data=stored_input,
        )
        if created:
            try:
                uploaded_inputs = self._save_uploads(run.id, uploads or {})
                execution_input = {**(input_data or {}), **uploaded_inputs}
                stored_input = self._redact_input(execution_input, sensitive_names)
                run = self.store.set_input(run.id, stored_input)
                if sensitive_names:
                    self._write_execution_input(run.id, execution_input)
            except (OSError, ValueError) as error:
                self._delete_uploads(run.id)
                self._delete_execution_input(run.id)
                return self.store.finish(
                    run.id,
                    status="failed",
                    message=f"Could not save uploaded file: {error}",
                ), True
            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)
            self._delete_execution_input(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,
        orphaned_task_state_retention_seconds: int = ORPHANED_TASK_STATE_RETENTION_SECONDS,
    ) -> None:
        """Remove transient run data and state for task groups removed long ago."""

        runs_root = self.settings.state_root / "runs"
        cutoff = time.time() - artifact_max_age_seconds
        if runs_root.is_dir():
            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)
            for path in runs_root.glob("*.execution-input.json"):
                run_id_text = path.name.removesuffix(".execution-input.json")
                if not run_id_text.isdigit():
                    continue
                run = self.store.get(int(run_id_text))
                if run is None or run.status not in {"queued", "running"}:
                    path.unlink(missing_ok=True)
        self._cleanup_orphaned_task_state(orphaned_task_state_retention_seconds)

    def _cleanup_orphaned_task_state(self, retention_seconds: int) -> None:
        """Delete state only after a group remains absent for a grace period."""

        if retention_seconds < 0 or not self.settings.task_root.is_dir():
            return
        root = self.settings.state_root / "tasks"
        if not root.is_dir():
            return
        marker_path = root / ORPHANED_TASK_STATE_FILENAME
        try:
            markers = json.loads(marker_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            markers = {}
        if not isinstance(markers, dict):
            markers = {}
        active_groups = {group.id for group in self.catalog.groups()}
        current_time = time.time()
        updated: dict[str, float] = {}
        for path in root.iterdir():
            if path.is_symlink() or not path.is_dir():
                continue
            if path.name in active_groups:
                continue
            first_missing_at = markers.get(path.name, current_time)
            if not isinstance(first_missing_at, (int, float)):
                first_missing_at = current_time
            if current_time - float(first_missing_at) >= retention_seconds:
                shutil.rmtree(path, ignore_errors=True)
            else:
                updated[path.name] = float(first_missing_at)
        if updated:
            temporary = marker_path.with_suffix(marker_path.suffix + ".tmp")
            temporary.write_text(json.dumps(updated, sort_keys=True) + "\n", encoding="utf-8")
            temporary.replace(marker_path)
        else:
            marker_path.unlink(missing_ok=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._execution_input_path(run.id)
        artifact_root = self._artifact_root(run.id)
        artifact_root.mkdir(parents=True, exist_ok=True)
        try:
            if not input_path.exists():
                input_path = self.settings.state_root / "runs" / f"{run.id}.input.json"
                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)
            self._delete_execution_input(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_ROOT": str(self.settings.state_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 _execution_input_path(self, run_id: int) -> Path:
        return self.settings.state_root / "runs" / f"{run_id}.execution-input.json"

    def _write_execution_input(
        self,
        run_id: int,
        input_data: Mapping[str, str | int | list[str]],
    ) -> None:
        path = self._execution_input_path(run_id)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(input_data, sort_keys=True, separators=(",", ":")) + "\n", encoding="utf-8")
        path.chmod(0o600)

    def _delete_execution_input(self, run_id: int) -> None:
        self._execution_input_path(run_id).unlink(missing_ok=True)

    @staticmethod
    def _redact_input(
        input_data: Mapping[str, str | int | list[str]],
        sensitive_names: set[str],
    ) -> dict[str, str | int | list[str]]:
        return {
            name: "[redacted]" if name in sensitive_names else value
            for name, value in input_data.items()
        }

    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