"""Start/stop/track workflow subprocesses."""

from __future__ import annotations

import asyncio
import os
import signal
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4

from mixersystem.studio.event_bus import EventBus, StudioEvent
from mixersystem.studio.session_paths import normalize_session_folder_arg


@dataclass
class WorkflowRun:
    run_id: str
    workflow: str
    session_folder: str
    args: dict = field(default_factory=dict)
    pid: int | None = None
    process: asyncio.subprocess.Process | None = field(default=None, repr=False)
    started_at: str = ""
    finished_at: str = ""
    exit_code: int | None = None
    status: str = "running"  # queued / running / completed / failed / killed
    queue_position: int | None = None  # 1-based position; None when not queued

    def to_dict(self) -> dict:
        return {
            "run_id": self.run_id,
            "workflow": self.workflow,
            "session_folder": self.session_folder,
            "args": self.args,
            "pid": self.pid,
            "started_at": self.started_at,
            "finished_at": self.finished_at,
            "exit_code": self.exit_code,
            "status": self.status,
            "queue_position": self.queue_position,
        }


class ProcessManager:
    """Manage workflow subprocess lifecycle."""

    def __init__(self, event_bus: EventBus, project_root: Path) -> None:
        self._event_bus = event_bus
        self._project_root = project_root
        self._runs: dict[str, WorkflowRun] = {}
        self._monitors: dict[str, asyncio.Task] = {}
        self._sessions_launching: set[str] = set()
        # Per-session queue: normalized session_folder -> ordered list of queued run_ids
        self._session_queues: dict[str, list[str]] = {}
        self._state_lock = asyncio.Lock()

    def get_run(self, run_id: str) -> WorkflowRun | None:
        return self._runs.get(run_id)

    def list_runs(self) -> list[WorkflowRun]:
        """Return active runs (running or queued only)."""
        return [r for r in self._runs.values() if r.status in ("running", "queued")]

    def list_active(self) -> list[WorkflowRun]:
        return [r for r in self._runs.values() if r.status == "running"]

    async def start_workflow(
        self, workflow: str, session_folder: str, args: dict | None = None,
    ) -> WorkflowRun:
        """Launch a workflow as a subprocess, or queue it if one is already running in this session."""
        args = args or {}
        run_id = uuid4().hex[:12]

        # Auto-generate session folder if empty.
        if not session_folder:
            import secrets

            session_folder = f"local-{secrets.token_hex(2)}"
        session_folder = normalize_session_folder_arg(session_folder)

        queued_run: WorkflowRun | None = None
        async with self._state_lock:
            if self._has_running_for_session(session_folder):
                queue = self._session_queues.setdefault(session_folder, [])
                queued_run = WorkflowRun(
                    run_id=run_id,
                    workflow=workflow,
                    session_folder=session_folder,
                    args=args,
                    status="queued",
                    queue_position=len(queue) + 1,
                )
                self._runs[run_id] = queued_run
                queue.append(run_id)
            else:
                # Reserve a launch slot so concurrent starts in the same session are queued.
                self._sessions_launching.add(session_folder)

        if queued_run:
            self._event_bus.emit(
                StudioEvent(
                    type="workflow_queued",
                    payload=queued_run.to_dict(),
                    run_id=run_id,
                    session_folder=session_folder,
                )
            )
            self._emit_runs_snapshot()
            return queued_run

        try:
            return await self._launch_run(run_id, workflow, session_folder, args)
        except Exception:
            async with self._state_lock:
                self._sessions_launching.discard(session_folder)
            raise

    async def stop_workflow(self, run_id: str) -> WorkflowRun | None:
        """Terminate a running or queued workflow."""
        proc: asyncio.subprocess.Process | None = None
        queued_stopped = False

        async with self._state_lock:
            run = self._runs.get(run_id)
            if not run:
                return None

            if run.status == "queued":
                queue = self._session_queues.get(run.session_folder, [])
                if run_id in queue:
                    queue.remove(run_id)
                if not queue:
                    self._session_queues.pop(run.session_folder, None)
                run.status = "killed"
                run.queue_position = None
                run.finished_at = datetime.now(timezone.utc).isoformat()
                self._recompute_queue_positions(run.session_folder)
                queued_stopped = True
            elif run.status == "running" and run.process:
                # Mark as killed before waiting so monitor path never emits workflow_finished.
                run.status = "killed"
                run.finished_at = datetime.now(timezone.utc).isoformat()
                proc = run.process
            else:
                return run

        if queued_stopped:
            self._event_bus.emit(
                StudioEvent(
                    type="workflow_stopped",
                    payload=run.to_dict(),
                    run_id=run_id,
                    session_folder=run.session_folder,
                )
            )
            self._emit_runs_snapshot()
            asyncio.get_running_loop().call_later(30, self._purge_run, run_id)
            return run

        assert proc is not None
        try:
            pgid = os.getpgid(proc.pid)
            os.killpg(pgid, signal.SIGTERM)
            try:
                await asyncio.wait_for(proc.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                os.killpg(pgid, signal.SIGKILL)
                await proc.wait()
        except (ProcessLookupError, OSError):
            try:
                await asyncio.wait_for(proc.wait(), timeout=2.0)
            except asyncio.TimeoutError:
                pass

        async with self._state_lock:
            run.exit_code = proc.returncode

        self._event_bus.stop_watching(run.session_folder, run_id)
        self._event_bus.emit(
            StudioEvent(
                type="workflow_stopped",
                payload=run.to_dict(),
                run_id=run_id,
                session_folder=run.session_folder,
            )
        )
        self._emit_runs_snapshot()

        await self._launch_queued(run.session_folder)
        asyncio.get_running_loop().call_later(30, self._purge_run, run_id)
        return run

    def _emit_process_output(self, run: WorkflowRun, stream: str, line: str) -> None:
        self._event_bus.emit(
            StudioEvent(
                type="process_output",
                payload={"stream": stream, "line": line},
                run_id=run.run_id,
                session_folder=run.session_folder,
            )
        )

    async def _drain_stream(
        self,
        run: WorkflowRun,
        stream_name: str,
        stream: asyncio.StreamReader | None,
    ) -> None:
        if stream is None:
            return

        pending = ""
        while True:
            chunk = await stream.read(4096)
            if not chunk:
                break
            pending += chunk.decode(errors="replace")
            while True:
                newline = pending.find("\n")
                if newline < 0:
                    break
                line = pending[:newline].rstrip("\r")
                pending = pending[newline + 1:]
                self._emit_process_output(run, stream_name, line)

        if pending:
            self._emit_process_output(run, stream_name, pending.rstrip("\r"))

    async def _monitor_process(self, run: WorkflowRun) -> None:
        """Drain process pipes and detect process exit."""
        proc = run.process
        if not proc:
            return

        drainers: list[asyncio.Task] = []
        if proc.stdout:
            drainers.append(asyncio.create_task(self._drain_stream(run, "stdout", proc.stdout)))
        if proc.stderr:
            drainers.append(asyncio.create_task(self._drain_stream(run, "stderr", proc.stderr)))

        try:
            if drainers:
                await asyncio.gather(*drainers)
        except (asyncio.CancelledError, OSError):
            for task in drainers:
                if not task.done():
                    task.cancel()

        await proc.wait()

        async with self._state_lock:
            if run.status != "running":
                should_emit_finished = False
            else:
                run.exit_code = proc.returncode
                run.finished_at = datetime.now(timezone.utc).isoformat()
                run.status = "completed" if proc.returncode == 0 else "failed"
                should_emit_finished = True

        if should_emit_finished:
            self._event_bus.stop_watching(run.session_folder, run.run_id)
            self._event_bus.emit(
                StudioEvent(
                    type="workflow_finished",
                    payload=run.to_dict(),
                    run_id=run.run_id,
                    session_folder=run.session_folder,
                )
            )
            self._emit_runs_snapshot()
            await self._launch_queued(run.session_folder)

        self._monitors.pop(run.run_id, None)
        # Clean up finished runs after a delay to avoid stale data in list_runs
        asyncio.get_running_loop().call_later(30, self._purge_run, run.run_id)

    async def _launch_run(
        self, run_id: str, workflow: str, session_folder: str, args: dict,
    ) -> WorkflowRun:
        """Spawn a subprocess for a workflow run."""
        cli_args = _build_cli_args(workflow, session_folder, args)

        proc = await asyncio.create_subprocess_exec(
            *cli_args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(self._project_root),
            start_new_session=True,
        )

        async with self._state_lock:
            run = self._runs.get(run_id)
            if run:
                # Promoting from queued -> running.
                run.pid = proc.pid
                run.process = proc
                run.started_at = datetime.now(timezone.utc).isoformat()
                run.status = "running"
                run.queue_position = None
            else:
                run = WorkflowRun(
                    run_id=run_id,
                    workflow=workflow,
                    session_folder=session_folder,
                    args=args,
                    pid=proc.pid,
                    process=proc,
                    started_at=datetime.now(timezone.utc).isoformat(),
                    status="running",
                )
                self._runs[run_id] = run
            self._sessions_launching.discard(session_folder)

        self._event_bus.watch_trace_log(session_folder, run_id)
        self._event_bus.emit(
            StudioEvent(
                type="workflow_started",
                payload=run.to_dict(),
                run_id=run_id,
                session_folder=session_folder,
            )
        )
        self._emit_runs_snapshot()

        monitor = asyncio.create_task(self._monitor_process(run))
        self._monitors[run_id] = monitor
        return run

    async def _launch_queued(self, session_folder: str) -> None:
        """Pop and launch the next queued run for a session, if any."""
        next_run_id = ""
        next_workflow = ""
        next_session_folder = session_folder
        next_args: dict = {}

        async with self._state_lock:
            if self._has_running_for_session(session_folder):
                return

            queue = self._session_queues.get(session_folder, [])
            while queue:
                candidate_run_id = queue.pop(0)
                run = self._runs.get(candidate_run_id)
                if not run or run.status != "queued":
                    continue
                next_run_id = candidate_run_id
                next_workflow = run.workflow
                next_session_folder = run.session_folder
                next_args = run.args
                self._sessions_launching.add(next_session_folder)
                break

            if not queue:
                self._session_queues.pop(session_folder, None)
            self._recompute_queue_positions(session_folder)

        if not next_run_id:
            return
        try:
            await self._launch_run(next_run_id, next_workflow, next_session_folder, next_args)
        except Exception:
            async with self._state_lock:
                self._sessions_launching.discard(next_session_folder)
                run = self._runs.get(next_run_id)
                if run and run.status == "queued":
                    run.status = "failed"
                    run.finished_at = datetime.now(timezone.utc).isoformat()
                    run.exit_code = 1
                else:
                    run = None
            if run:
                self._event_bus.emit(
                    StudioEvent(
                        type="workflow_finished",
                        payload=run.to_dict(),
                        run_id=next_run_id,
                        session_folder=next_session_folder,
                    )
                )
            await self._launch_queued(next_session_folder)

    def _recompute_queue_positions(self, session_folder: str) -> None:
        """Update queue_position on remaining queued runs and emit events."""
        queue = self._session_queues.get(session_folder, [])
        for i, rid in enumerate(queue):
            run = self._runs.get(rid)
            if run and run.status == "queued":
                run.queue_position = i + 1
                self._event_bus.emit(
                    StudioEvent(
                        type="queue_updated",
                        payload={"run_id": rid, "queue_position": run.queue_position},
                        run_id=rid,
                        session_folder=session_folder,
                    )
                )

    def _emit_runs_snapshot(self) -> None:
        """Broadcast the full list of active runs so the frontend stays in sync."""
        runs = [r.to_dict() for r in self._runs.values()
                if r.status in ("running", "queued")]
        self._event_bus.emit(
            StudioEvent(type="runs_snapshot", payload={"runs": runs}))

    def _purge_run(self, run_id: str) -> None:
        """Remove a finished run from memory."""
        run = self._runs.get(run_id)
        if run and run.status not in ("running", "queued"):
            del self._runs[run_id]

    def _has_running_for_session(self, session_folder: str) -> bool:
        if session_folder in self._sessions_launching:
            return True
        return any(
            r.session_folder == session_folder and r.status == "running"
            for r in self._runs.values()
        )

    async def shutdown(self) -> None:
        """Kill all running processes on server stop."""
        for run in list(self._runs.values()):
            if run.status in ("running", "queued"):
                await self.stop_workflow(run.run_id)
        for task in self._monitors.values():
            if not task.done():
                task.cancel()
        if self._monitors:
            await asyncio.gather(*self._monitors.values(), return_exceptions=True)


def _build_cli_args(workflow: str, session_folder: str, args: dict) -> list[str]:
    """Build CLI command list for spawning a workflow subprocess."""
    cmd = [
        sys.executable,
        "-m",
        "mixersystem",
        "run",
        workflow,
        f"--session_folder={session_folder}",
    ]
    for key, value in args.items():
        if isinstance(value, bool):
            if value:
                cmd.append(f"--{key}=true")
        elif value is not None and str(value):
            cmd.append(f"--{key}={value}")
    return cmd
