from typing import Optional, List, Union, Protocol, runtime_checkable
from pydantic import Field, validator
from koil.composition import KoiledModel
from pathlib import Path
import asyncio
from datetime import timedelta
from typing import Union, Dict, Literal
from .compose_spec import ComposeSpec
import json
import os

ValidPath = Union[str, Path]


@runtime_checkable
class CLIBearer(Protocol):
    async def aget_cli(self) -> "CLI":
        ...


class CLI(KoiledModel):
    """A CLI object that represents the docker-compose CLI.

    This is a pydantic model that can be used to build the docker-compose CLI
    command. It also contains methods for running the CLI command
    asynchronously.

    """

    config: Optional[ValidPath] = None
    context: Optional[str] = None
    debug: Optional[bool] = None
    host: Optional[str] = None
    log_level: Optional[str] = None
    tls: Optional[bool] = None
    tlscacert: Optional[ValidPath] = None
    tlscert: Optional[ValidPath] = None
    tlskey: Optional[ValidPath] = None
    tlsverify: Optional[bool] = None
    compose_files: List[ValidPath] = Field(
        default_factory=lambda: ["docker-compose.yml"]
    )
    compose_profiles: List[ValidPath] = Field(default_factory=list)
    compose_env_file: Optional[ValidPath] = Field(default=".env")
    compose_project_name: Optional[str] = None
    compose_project_directory: Optional[ValidPath] = None
    compose_compatibility: Optional[bool] = None
    client_call: List[str] = Field(default_factory=lambda: ["docker", "compose"])

    @validator("compose_files", each_item=True)
    def _validate_compose_files(cls, v):
        if os.path.exists(v):
            return v
        else:
            raise ValueError(f"Compose File {v} does not exist.")

    @property
    def docker_cmd(self) -> list[str]:
        """Builds the docker command. This is the base prepended
        command that will be run by the CLI.
        """
        result = self.client_call

        if self.compose_files:
            for compose_file in self.compose_files:
                result += ["--file", compose_file]

        if self.config is not None:
            result += ["--config", self.config]

        if self.context is not None:
            result += ["--context", self.context]

        if self.debug:
            result.append("--debug")

        if self.host is not None:
            result += ["--host", self.host]

        if self.log_level is not None:
            result += ["--log-level", self.log_level]

        if self.tls:
            result.append("--tls")

        if self.tlscacert is not None:
            result += ["--tlscacert", self.tlscacert]

        if self.tlscert is not None:
            result += ["--tlscert", self.tlscert]

        if self.tlskey is not None:
            result += ["--tlskey", self.tlskey]

        if self.tlsverify:
            result.append("--tlsverify")

        return result

    async def arun(command: str):
        proc = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        stdout, stderr = await proc.communicate()

        return stdout.decode("utf-8"), stderr.decode("utf-8")

    async def astream_docker_logs(
        self,
        tail: Optional[str] = None,
        follow: bool = False,
        no_log_prefix: bool = False,
        timestamps: bool = False,
        since: Optional[str] = None,
        until: Optional[str] = None,
        services: Union[str, List[str]] = [],
    ):
        full_cmd = self.docker_cmd + ["logs", "--no-color"]
        if tail is not None:
            full_cmd += ["--tail", tail]
        if follow:
            full_cmd.append("--follow")
        if no_log_prefix:
            full_cmd.append("--no-log-prefix")
        if timestamps:
            full_cmd.append("--timestamps")
        if since is not None:
            full_cmd += ["--since", since]
        if until is not None:
            full_cmd += ["--until", until]

        if services:
            if isinstance(services, str):
                services = [services]
            full_cmd += services

        async for line in self.astream_command(full_cmd):
            yield line

    async def _astread_stream(
        self,
        stream: asyncio.StreamReader,
        queue: asyncio.Queue,
        name: str,
    ):
        async for line in stream:
            await queue.put((name, line.decode("utf-8").strip()))

        await queue.put(None)

    async def astream_command(self, command: List[str]):
        # Create the subprocess using asyncio's subprocess

        full_cmd = " ".join(map(str, command))

        proc = await asyncio.create_subprocess_shell(
            full_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        queue = asyncio.Queue()
        # Create and start tasks for reading each stream

        try:
            readers = [
                asyncio.create_task(self._astread_stream(proc.stdout, queue, "STDOUT")),
                asyncio.create_task(self._astread_stream(proc.stderr, queue, "STDERR")),
            ]

            # Track the number of readers that are finished
            finished_readers = 0
            while finished_readers < len(readers):
                line = await queue.get()
                if line is None:
                    finished_readers += 1  # One reader has finished
                    continue
                yield line

            # Cleanup: cancel any remaining reader tasks
            for reader in readers:
                reader.cancel()
                try:
                    await reader
                except asyncio.CancelledError:
                    pass

        except asyncio.CancelledError:
            # Handle cancellation request
            proc.kill()
            await proc.wait()  # Wait for the subprocess to exit after receiving SIGINT

            # Cleanup: cancel any remaining reader tasks
            for reader in readers:
                reader.cancel()
                try:
                    await reader
                except asyncio.CancelledError:
                    pass

            raise

        except Exception as e:
            raise e

    async def astream_pull(
        self,
        services: Union[List[str], str, None] = None,
        ignore_pull_failures: bool = False,
        include_deps: bool = False,
        quiet: bool = False,
    ):
        full_cmd = self.docker_cmd + ["pull"]
        if ignore_pull_failures:
            full_cmd.append("--ignore-pull-failures")
        if include_deps:
            full_cmd.append("--include-deps")
        if quiet:
            full_cmd.append("--quiet")

        if services:
            if isinstance(services, str):
                services = [services]
            full_cmd += services

        async for line in self.astream_command(full_cmd):
            yield line

    async def astream_stop(
        self,
        services: Union[str, List[str], None] = None,
        timeout: Union[int, timedelta, None] = None,
        stream_logs: bool = False,
    ):
        full_cmd = self.docker_cmd + ["stop"]
        if timeout is not None:
            if isinstance(timeout, timedelta):
                timeout = int(timeout.total_seconds())

            full_cmd.append(f"--timeout {timeout}")

        if services:
            if isinstance(services, str):
                services = [services]
            full_cmd += services

        async for line in self.astream_command(full_cmd):
            yield line

    async def astream_restart(
        self,
        services: Union[str, List[str], None] = None,
    ):
        full_cmd = self.docker_cmd + ["restart"]

        if services:
            if isinstance(services, str):
                services = [services]
            full_cmd += services

        async for line in self.astream_command(full_cmd):
            yield line

    async def astream_up(
        self,
        services: Union[List[str], str, None] = None,
        build: bool = False,
        detach: bool = False,
        abort_on_container_exit: bool = False,
        scales: Dict[str, int] = {},
        attach_dependencies: bool = False,
        force_recreate: bool = False,
        no_recreate: bool = False,
        no_build: bool = False,
        remove_orphans: bool = False,
        renew_anon_volumes: bool = False,
        no_color: bool = False,
        no_log_prefix: bool = False,
        no_start: bool = False,
        quiet: bool = False,
        wait: bool = False,
        no_attach_services: Union[List[str], str, None] = None,
        pull: Literal["always", "missing", "never", None] = None,
        stream_logs: bool = False,
    ):
        if quiet and stream_logs:
            raise ValueError(
                "It's not possible to have stream_logs=True and quiet=True at the same time. "
                "Only one can be activated at a time."
            )
        full_cmd = self.docker_cmd + ["up"]
        if build:
            full_cmd.append("--build")
        if detach:
            full_cmd.append("--detach")
        if abort_on_container_exit:
            full_cmd.append("--abort-on-container-exit")
        for service, scale in scales.items():
            full_cmd.append(f"--scale {service}={scale}")
        if attach_dependencies:
            full_cmd.append("--attach-dependencies")
        if force_recreate:
            full_cmd.append("--force-recreate")
        if no_recreate:
            full_cmd.append("--no-recreate")
        if no_build:
            full_cmd.append("--no-build")
        if remove_orphans:
            full_cmd.append("--remove-orphans")
        if renew_anon_volumes:
            full_cmd.append("--renew-anon-volumes")
        if no_color:
            full_cmd.append("--no-color")
        if no_log_prefix:
            full_cmd.append("--no-log-prefix")
        if no_start:
            full_cmd.append("--no-start")
        if quiet:
            full_cmd.append("--quiet")
        if wait:
            full_cmd.append("--wait")
        if no_attach_services is not None:
            if isinstance(no_attach_services, str):
                no_attach_services = [no_attach_services]
            for service in no_attach_services:
                full_cmd.append(f"--no-attach {service}")
        if pull is not None:
            full_cmd.append(f"--pull {pull}")

        if services:
            if isinstance(services, str):
                services = [services]
            full_cmd += services

        async for line in self.astream_command(full_cmd):
            yield line

    async def ainspect_config(self) -> ComposeSpec:
        """Returns the configuration of the compose stack for further inspection.

        For example
        ```python
        from python_on_whales import docker
        project_config = docker.compose.config()
        print(project_config.services["my_first_service"].image)
        "redis"
        ```

        Parameters:
            return_json: If `False`, a `ComposeConfig` object will be returned, and you
                'll be able to take advantage of your IDE autocompletion. If you want the
                full json output, you may use `return_json`. In this case, you'll get
                lists and dicts corresponding to the json response, unmodified.
                It may be useful if you just want to print the config or want to access
                a field that was not in the `ComposeConfig` class.

        # Returns
            A `ComposeConfig` object if `return_json` is `False`, and a `dict` otherwise.
        """
        full_cmd = self.docker_cmd + ["config", "--format", "json"]

        lines = []

        async for x, line in self.astream_command(full_cmd):
            lines.append(line)

        result = "\n".join(lines)

        try:
            return ComposeSpec(**json.loads(result))
        except Exception as e:
            raise Exception(
                f"Could not inspect! Error while parsing the json: {result}"
            ) from e
