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

from __future__ import annotations

import asyncio
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 list(self._runs.values())

    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,
                )
            )
            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,
                )
            )
            return run

        assert proc is not None
        try:
            proc.terminate()
            try:
                await asyncio.wait_for(proc.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                proc.kill()
                await proc.wait()
        except ProcessLookupError:
            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,
            )
        )

        await self._launch_queued(run.session_folder)
        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,
                )
            )
            await self._launch_queued(run.session_folder)

        self._monitors.pop(run.run_id, None)

    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),
        )

        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,
            )
        )

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