from typing import Optional, TypedDict
from aiohttp import web, WSMsgType, ClientSession
from collections import defaultdict
import asyncio
import logging
import os
import json
from asyncio.subprocess import Process
import psutil


logger = logging.getLogger(__name__)

PROCESS_OWNERSHIP: dict[int, Process] = {}
SUBSCRIPTIONS: defaultdict[int, list[web.WebSocketResponse]] = defaultdict(list)
PROCESS_OWNERSHIP_LOCK = asyncio.Lock()
SUBSCRIPTIONS_LOCK = asyncio.Lock()


class SpawnProcessRequest(TypedDict):
    cmd: str
    args: list[str]
    env: dict[str, str]


class StopProcessRequest(TypedDict):
    pid: int


async def check_processes_step():
    print(PROCESS_OWNERSHIP)
    for pid, process in list(PROCESS_OWNERSHIP.items()):
        try:
            if (
                psutil.pid_exists(pid)
                and psutil.Process(pid).status() == psutil.STATUS_RUNNING
            ):
                continue
        except psutil.NoSuchProcess:
            pass

        logger.info(f"Process {pid} is not running ({process.returncode})")
        await stop_subprocess(process, pid)

    async with SUBSCRIPTIONS_LOCK:
        for pid, subs in list(SUBSCRIPTIONS.items()):
            if pid not in PROCESS_OWNERSHIP:
                for ws in subs:
                    await ws.close()
                del SUBSCRIPTIONS[pid]


async def run_subprocess_monitor(port: int = 5057, check_interval: float = 2):
    check_interval = max(0.1, check_interval)  # Ensure period is not too small
    app = web.Application()

    # the index page shows the current status of the subprocesses
    async def index(req):
        return web.json_response(list(PROCESS_OWNERSHIP.keys()))

    async def spawn(req: web.Request):
        request: SpawnProcessRequest = await req.json()
        # to avoid thread safety issues as the web framework used here is not mandatory
        try:
            subprocess_pid = await start_subprocess(request, port)
            return web.json_response({"code": "success", "pid": subprocess_pid})
        except Exception as exc:
            cmd = " ".join([request["cmd"], *request["args"]])
            logger.error(f"Failed to start subprocess: {cmd}")
            logger.exception(exc)
            return web.json_response({"code": "failure", "error": str(exc)})

    async def stop(req: web.Request):
        request: StopProcessRequest = await req.json()
        found = await stop_subprocess_request(request, asyncio.get_running_loop())
        return web.json_response({"code": "success" if found else "failure"})

    app.router.add_get("/", index)
    app.router.add_post("/spawn", spawn)
    app.router.add_post("/stop", stop)
    app.router.add_get("/subscribe", subscribe_output)  # New endpoint for subscriptions

    async def serve():
        _scan_period = check_interval
        runner = web.AppRunner(app)

        try:
            logger.info(f"Starting subprocess manager on port {port}...")
            await runner.setup()
            site = web.TCPSite(runner, "localhost", port)
            await site.start()

            while True:
                await asyncio.sleep(_scan_period)
                await check_processes_step()

        finally:
            await runner.cleanup()
            await kill_all_subprocesses()

            # Close all WebSocket connections
            async with SUBSCRIPTIONS_LOCK:
                for subs in SUBSCRIPTIONS.values():
                    for ws in subs:
                        await ws.close()

    await serve()


async def subscribe_output(request: web.Request):
    pid = int(request.query.get("pid", -1))
    if pid == -1 or pid not in PROCESS_OWNERSHIP:
        return web.HTTPBadRequest(text="Invalid or missing 'pid' parameter.")

    ws = web.WebSocketResponse()
    await ws.prepare(request)
    async with SUBSCRIPTIONS_LOCK:
        SUBSCRIPTIONS[pid].append(ws)
    logger.info(f"Client subscribed to subprocess {pid} output.")

    try:
        async for msg in ws:
            if msg.type == WSMsgType.TEXT:
                # You can handle incoming messages from the client here if needed
                pass
            elif msg.type == WSMsgType.ERROR:
                logger.exception(ws.exception())
                logger.error(
                    f"WebSocket connection closed with exception {ws.exception()}"
                )
    finally:
        async with SUBSCRIPTIONS_LOCK:
            if pid in SUBSCRIPTIONS and ws in SUBSCRIPTIONS[pid]:
                SUBSCRIPTIONS[pid].remove(ws)
        logger.info(f"Client unsubscribed from subprocess {pid} output.")

    return ws


def remote_spawn_subprocess(
    command: str, args: list[str], env: dict[str, str], port: int = 5687
):
    """
    sends a spwan request to the service

    command: the command to spawn
    args: the arguments of the command
    env: the environment variables
    port: the port that the service is deployed on
    """

    async def send_request():
        req = SpawnProcessRequest(cmd=command, args=args, env=env)
        logger.info(f"Sending request to spawn subprocess: {json.dumps(req, indent=2)}")
        async with ClientSession() as session:
            async with session.post(
                f"http://localhost:{port}/spawn",
                json=req,
            ) as resp:
                ans = await resp.json()
                logger.info(json.dumps(ans, indent=2, ensure_ascii=True))
                return ans

    return asyncio.run(send_request())


async def kill_all_subprocesses():
    logger.info("Killing all subprocesses...")
    for pid, process in list(PROCESS_OWNERSHIP.items()):
        await stop_subprocess(process, pid)


def terminate_subprocess_sync(process: Process):
    try:
        process.terminate()
        logger.info(f"Terminated subprocess {process.pid}")
    except ProcessLookupError:
        pass
    except Exception as exc:
        logger.exception(exc)
        logger.error(f"Error terminating subprocess {process.pid}: {exc}")


def kill_subprocess_sync(process: Process):
    try:
        process.kill()
        logger.warning(f"Killed subprocess {process.pid}")
    except Exception as exc:
        logger.exception(exc)
        logger.error(f"Error killing subprocess {process.pid}: {exc}")


async def stop_subprocess(
    process: Optional[Process] = None,
    pid: Optional[int] = None,
    loop: Optional[asyncio.AbstractEventLoop] = None,
):
    if loop is None:
        loop = asyncio.get_running_loop()
    if process is None:
        if pid is None:
            raise ValueError("Either process or pid must be provided")

        if pid not in PROCESS_OWNERSHIP:
            raise ValueError("PID not found")
        if PROCESS_OWNERSHIP[pid].pid == pid:
            process = PROCESS_OWNERSHIP[pid]

    if process is None:
        raise ValueError("Process not found")

    if pid is None:
        for _pid, process in PROCESS_OWNERSHIP.items():
            if process == process:
                pid = _pid
                break

    if pid is None:
        raise ValueError("PID not found")

    terminate_subprocess_sync(process)

    async def check_terminated(p: Process, pid):
        try:
            await p.wait()
        except Exception:
            pass
        if p.returncode is None:
            kill_subprocess_sync(p)

        async with PROCESS_OWNERSHIP_LOCK:
            del PROCESS_OWNERSHIP[pid]

    loop.call_soon_threadsafe(asyncio.create_task, check_terminated(process, pid))


async def stop_subprocess_request(
    request: StopProcessRequest, loop: Optional[asyncio.AbstractEventLoop] = None
):
    if loop is None:
        loop = asyncio.get_running_loop()

    pid = request["pid"]
    if pid not in PROCESS_OWNERSHIP:
        return False
    else:
        if PROCESS_OWNERSHIP[pid].pid == pid:
            await stop_subprocess(PROCESS_OWNERSHIP[pid], pid, loop)
            return True
    return False


async def broadcast_output(pid: int, message: str):
    subscribers = SUBSCRIPTIONS.get(pid, [])
    await asyncio.gather(*[ws.send_str(message) for ws in subscribers if not ws.closed])


async def stream_subprocess_output(pid: int, process: Process):
    async def read_stream(stream, stream_name):
        while True:
            line = await stream.readline()
            if line:
                message = json.dumps(
                    {
                        "pid": process.pid,
                        "stream": stream_name,
                        "data": line.decode().rstrip(),
                    }
                )
                await broadcast_output(pid, message)
            else:
                break

    await asyncio.gather(
        read_stream(process.stdout, "stdout"), read_stream(process.stderr, "stderr")
    )


async def start_subprocess(request: SpawnProcessRequest, port: int):
    cmd = request["cmd"]
    args = request["args"]
    env = request.get("env", {})

    logger.info(f"Starting subprocess: {cmd} {args} with environment: {env}")
    full_command = [cmd] + args

    env["SUBPROCESS_MONITOR_PORT"] = str(port)

    env = {**os.environ, **env}

    process = await asyncio.create_subprocess_exec(
        *full_command,
        env=env,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        # stdin=asyncio.subprocess.PIPE,
    )

    async with PROCESS_OWNERSHIP_LOCK:
        PROCESS_OWNERSHIP[process.pid] = process
    logger.info(f"Started subprocess: {cmd} {' '.join(args)} with PID {process.pid}")
    # Start tasks to read stdout and stderr
    asyncio.create_task(stream_subprocess_output(process.pid, process))
    return process.pid
