from typing import Literal

import anyio
import anyio.abc
from aiodocker import Docker, containers
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from loguru import logger
from pydantic import BaseModel, Field


class DockerServerParameters(BaseModel):
    """Configuration for Docker transport."""

    image: str = Field(description="Image of the docker container")
    args: list[str] = Field(
        default_factory=list,
        description="Command line arguments for the docker container",
    )
    env: dict[str, str] = Field(
        default_factory=dict,
        description="Environment variables for the docker container",
    )


class DockerTransport:
    state: Literal["constructed", "initialized", "started", "stopped"] = "constructed"

    def __init__(self, config: DockerServerParameters) -> None:
        self.config = config
        self.docker: Docker | None = None
        self.container: containers.DockerContainer | None = None
        self.attach_result: any = None

        self._reader_send: MemoryObjectSendStream[str]
        self._reader_recv: MemoryObjectReceiveStream[str]

        self._writer_send: MemoryObjectSendStream[str]
        self._writer_recv: MemoryObjectReceiveStream[str]

        self._task_group: anyio.abc.TaskGroup | None = None

    async def init(self) -> None:
        self.docker = Docker()
        await self.docker.images.pull(self.config.image)

        self._reader_send, self._reader_recv = anyio.create_memory_object_stream()
        self._writer_send, self._writer_recv = anyio.create_memory_object_stream()

        self.state = "initialized"
        logger.debug("DockerTransport initialized")

    async def start(self) -> None:
        if not self.docker:
            raise RuntimeError("Docker not initialized")

        self.container = await self.docker.containers.create(
            {
                "Image": self.config.image,
                "Cmd": self.config.args,
                "OpenStdin": True,
                "AttachStdin": True,
                "AttachStdout": True,
                "AttachStderr": True,
                "Tty": False,
                "Env": [f"{k}={v}" for k, v in self.config.env.items()],
                "HostConfig": {"AutoRemove": True},
            }
        )

        await self.container.start()
        logger.debug(f"Started container {self.container.id}")

        self.attach_result = self.container.attach(stdin=True, stdout=True, stderr=True)

        async def read_stdout():
            assert self.attach_result is not None

            buffer = ""
            async with self._reader_send:
                while True:
                    try:
                        msg = await self.attach_result.read_out()
                        if msg is None:
                            logger.debug("Docker stream closed (EOF)")
                            break
                        if msg.data is None:
                            continue

                        if isinstance(msg.data, (bytes, bytearray)):
                            chunk = msg.data.decode("utf-8", errors="replace")
                        else:
                            chunk = str(msg.data)

                        buffer += chunk
                        while "\n" in buffer:
                            line, buffer = buffer.split("\n", 1)
                            await self._reader_send.send(line)
                    except Exception as e:
                        logger.error(f"Error reading container output: {e}")
                        break
            logger.debug("read_stdout task exited")

        async def write_stdin():
            assert self.attach_result is not None
            async with self._writer_recv:
                async for msg in self._writer_recv:
                    try:
                        await self.attach_result.write_in(msg.encode("utf-8") + b"\n")
                    except Exception as e:
                        logger.warning(f"Failed to write to container stdin: {e}")
            logger.debug("write_stdin task exited")

        self._task_group = anyio.create_task_group()
        await self._task_group.__aenter__()
        self._task_group.start_soon(read_stdout)
        self._task_group.start_soon(write_stdin)

        self.state = "started"

    async def stop(self) -> None:
        logger.debug("Stopping DockerTransport...")

        if self.attach_result:
            try:
                await self.attach_result.close()
            except Exception:
                pass
            self.attach_result = None

        if self.container:
            try:
                await self.container.stop()
            except Exception:
                pass
            try:
                await self.container.delete()
            except Exception:
                pass
            self.container = None

        # Close input stream to stop write_stdin task
        try:
            await self._writer_send.aclose()
        except Exception:
            pass

        if self._task_group:
            await self._task_group.__aexit__(None, None, None)
            self._task_group = None

        if self.docker:
            await self.docker.close()
            self.docker = None

        self.state = "stopped"
        logger.debug("DockerTransport stopped and cleaned up")


    async def send(self, message: str) -> None:
        if self.state != "started":
            raise RuntimeError("Transport not started")
        await self._writer_send.send(message)

    async def receive(self) -> str:
        if self.state != "started":
            raise RuntimeError("Transport not started")
        return await self._reader_recv.receive()
