from __future__ import annotations

import asyncio
import os
from collections import OrderedDict
from collections.abc import Mapping, MutableMapping, Sequence
from typing import Any, override

from ai.backend.agent.kernel import AbstractCodeRunner, AbstractKernel, NextResult, ResultRecord
from ai.backend.agent.resources import KernelResourceSpec
from ai.backend.agent.types import AgentEventData, KernelOwnershipData
from ai.backend.common.docker import ImageRef
from ai.backend.common.dto.agent.response import CodeCompletionResp, CodeCompletionResult
from ai.backend.common.events.dispatcher import EventProducer
from ai.backend.common.types import CommitStatus, KernelId, SessionId


class DummyKernel(AbstractKernel):
    dummy_config: Mapping[str, Any]

    def __init__(
        self,
        ownership_data: KernelOwnershipData,
        network_id: str,
        image: ImageRef,
        version: int,
        *,
        agent_config: Mapping[str, Any],
        resource_spec: KernelResourceSpec,
        service_ports: Any,  # TODO: type-annotation
        environ: Mapping[str, Any],
        data: dict[str, Any],
        dummy_config: Mapping[str, Any],
    ) -> None:
        super().__init__(
            ownership_data,
            network_id,
            image,
            version,
            agent_config=agent_config,
            resource_spec=resource_spec,
            service_ports=service_ports,
            data=data,
            environ=environ,
        )
        self.is_commiting = False
        self.dummy_config = dummy_config
        self.dummy_kernel_cfg = self.dummy_config["kernel"]

    @override
    async def close(self) -> None:
        pass

    @override
    async def create_code_runner(
        self,
        event_producer: EventProducer,
        *,
        client_features: frozenset[str],
        api_version: int,
    ) -> AbstractCodeRunner:
        if self.dummy_kernel_cfg["use-fake-code-runner"]:
            return await DummyFakeCodeRunner.new(
                self.kernel_id,
                self.session_id,
                event_producer,
                kernel_host=self.data["kernel_host"],
                repl_in_port=self.data["repl_in_port"],
                repl_out_port=self.data["repl_out_port"],
                exec_timeout=0,
                client_features=client_features,
            )
        return await DummyCodeRunner.new(
            self.kernel_id,
            self.session_id,
            event_producer,
            kernel_host=self.data["kernel_host"],
            repl_in_port=self.data["repl_in_port"],
            repl_out_port=self.data["repl_out_port"],
            exec_timeout=0,
            client_features=client_features,
        )

    @override
    async def check_status(self) -> dict[str, Any]:
        delay = self.dummy_kernel_cfg["delay"]["check-status"]
        await asyncio.sleep(delay)
        return {}

    @override
    async def get_completions(self, text: str, opts: Mapping[str, Any]) -> CodeCompletionResp:
        delay = self.dummy_kernel_cfg["delay"]["get-completions"]
        await asyncio.sleep(delay)
        return CodeCompletionResp(result=CodeCompletionResult.success({"suggestions": []}))

    @override
    async def get_logs(self) -> dict[str, Any]:
        delay = self.dummy_kernel_cfg["delay"]["get-logs"]
        await asyncio.sleep(delay)
        return {"logs": "my logs"}

    @override
    async def interrupt_kernel(self) -> dict[str, Any]:
        delay = self.dummy_kernel_cfg["delay"]["interrupt-kernel"]
        await asyncio.sleep(delay)
        return {}

    @override
    async def start_service(self, service: str, opts: Mapping[str, Any]) -> dict[str, Any]:
        delay = self.dummy_kernel_cfg["delay"]["start-service"]
        await asyncio.sleep(delay)
        return {}

    @override
    async def start_model_service(self, model_service: Mapping[str, Any]) -> dict[str, Any]:
        delay = self.dummy_kernel_cfg["delay"]["start-model-service"]
        await asyncio.sleep(delay)
        return {}

    @override
    async def shutdown_service(self, service: str) -> None:
        delay = self.dummy_kernel_cfg["delay"]["shutdown-service"]
        await asyncio.sleep(delay)

    @override
    async def check_duplicate_commit(self, kernel_id: KernelId, subdir: str) -> CommitStatus:
        if self.is_commiting:
            return CommitStatus.ONGOING
        return CommitStatus.READY

    @override
    async def commit(
        self,
        kernel_id: KernelId,
        subdir: str,
        *,
        canonical: str | None = None,
        filename: str | None = None,
        extra_labels: dict[str, str] | None = None,
    ) -> None:
        if extra_labels is None:
            extra_labels = {}
        self.is_commiting = True
        delay = self.dummy_kernel_cfg["delay"]["commit"]
        await asyncio.sleep(delay)
        self.is_commiting = False

    @override
    async def get_service_apps(self) -> dict[str, Any]:
        delay = self.dummy_kernel_cfg["delay"]["get-service-apps"]
        await asyncio.sleep(delay)
        return {
            "status": "done",
            "data": [],
        }

    @override
    async def accept_file(self, container_path: os.PathLike[str] | str, filedata: bytes) -> None:
        delay = self.dummy_kernel_cfg["delay"]["accept-file"]
        await asyncio.sleep(delay)

    @override
    async def download_file(self, container_path: os.PathLike[str] | str) -> bytes:
        delay = self.dummy_kernel_cfg["delay"]["download-file"]
        await asyncio.sleep(delay)
        return b""

    @override
    async def download_single(self, container_path: os.PathLike[str] | str) -> bytes:
        delay = self.dummy_kernel_cfg["delay"]["download-single"]
        await asyncio.sleep(delay)
        return b""

    @override
    async def list_files(self, container_path: os.PathLike[str] | str) -> dict[str, Any]:
        delay = self.dummy_kernel_cfg["delay"]["list-files"]
        await asyncio.sleep(delay)
        return {"files": "", "errors": "", "abspath": ""}

    @override
    async def notify_event(self, evdata: AgentEventData) -> None:
        raise NotImplementedError


class DummyCodeRunner(AbstractCodeRunner):
    kernel_host: str
    repl_in_port: int
    repl_out_port: int

    def __init__(
        self,
        kernel_id: KernelId,
        session_id: SessionId,
        event_producer: EventProducer,
        *,
        kernel_host: str,
        repl_in_port: int,
        repl_out_port: int,
        exec_timeout: float = 0,
        client_features: frozenset[str] | None = None,
    ) -> None:
        super().__init__(
            kernel_id,
            session_id,
            event_producer,
            exec_timeout=exec_timeout,
            client_features=client_features,
        )
        self.kernel_host = kernel_host
        self.repl_in_port = repl_in_port
        self.repl_out_port = repl_out_port

    @override
    async def get_repl_in_addr(self) -> str:
        return f"tcp://{self.kernel_host}:{self.repl_in_port}"

    @override
    async def get_repl_out_addr(self) -> str:
        return f"tcp://{self.kernel_host}:{self.repl_out_port}"


class DummyFakeCodeRunner(AbstractCodeRunner):
    kernel_host: str
    repl_in_port: int
    repl_out_port: int

    input_sock: None
    output_sock: None
    zctx: None  # type: ignore[assignment]

    def __init__(
        self,
        kernel_id: KernelId,
        session_id: SessionId,
        event_producer: EventProducer,
        *,
        kernel_host: str,
        repl_in_port: int,
        repl_out_port: int,
        exec_timeout: float = 0,
        client_features: frozenset[str] | None = None,
    ) -> None:
        self.zctx = None
        self.input_sock = None
        self.output_sock = None

        self.completion_queue = asyncio.Queue(maxsize=128)
        self.service_queue = asyncio.Queue(maxsize=128)
        self.model_service_queue = asyncio.Queue(maxsize=128)
        self.service_apps_info_queue = asyncio.Queue(maxsize=128)
        self.status_queue = asyncio.Queue(maxsize=128)
        self.output_queue = None
        self.pending_queues = OrderedDict()
        self.current_run_id = None
        self.read_task = None
        self.status_task = None
        self.watchdog_task = None
        self._closed = False

        self.kernel_host = kernel_host
        self.repl_in_port = repl_in_port
        self.repl_out_port = repl_out_port

        self.event_producer = event_producer

    async def __ainit__(self) -> None:
        return

    def __setstate__(self, props: MutableMapping[str, Any]) -> None:
        self.__dict__.update(props)
        self.zctx = None
        self.input_sock = None
        self.output_sock = None

        self.completion_queue = asyncio.Queue(maxsize=128)
        self.service_queue = asyncio.Queue(maxsize=128)
        self.model_service_queue = asyncio.Queue(maxsize=128)
        self.service_apps_info_queue = asyncio.Queue(maxsize=128)
        self.status_queue = asyncio.Queue(maxsize=128)
        self.output_queue = None
        self.pending_queues = OrderedDict()
        self.current_run_id = None
        self.read_task = None
        self.status_task = None
        self.watchdog_task = None
        self._closed = False

    @override
    async def get_repl_in_addr(self) -> str:
        return f"tcp://{self.kernel_host}:{self.repl_in_port}"

    @override
    async def get_repl_out_addr(self) -> str:
        return f"tcp://{self.kernel_host}:{self.repl_out_port}"

    @override
    async def close(self) -> None:
        return None

    @override
    async def ping_status(self) -> None:
        return None

    @override
    async def feed_batch(self, opts: Mapping[str, Any]) -> None:
        return None

    @override
    async def feed_code(self, text: str) -> None:
        return None

    @override
    async def feed_input(self, text: str) -> None:
        return None

    @override
    async def feed_interrupt(self) -> None:
        return None

    @override
    async def feed_and_get_status(self) -> None:
        return None

    @override
    async def feed_and_get_completion(
        self, code_text: str, opts: Mapping[str, Any]
    ) -> CodeCompletionResult:
        return CodeCompletionResult.failure("not-implemented")

    @override
    async def feed_start_model_service(self, model_info: Mapping[str, Any]) -> dict[str, Any]:
        return {"status": "failed", "error": "not-implemented"}

    @override
    async def feed_start_service(self, service_info: Mapping[str, Any]) -> dict[str, Any]:
        return {"status": "failed", "error": "not-implemented"}

    @override
    async def feed_service_apps(self) -> dict[str, Any]:
        return {"status": "failed", "error": "not-implemented"}

    @override
    @staticmethod
    def aggregate_console(
        result: NextResult, records: Sequence[ResultRecord], api_ver: int
    ) -> None:
        return

    @override
    async def get_next_result(self, api_ver: int = 2, flush_timeout: float = 2.0) -> NextResult:
        return {
            "runId": self.current_run_id,
            "status": "finished",
            "exitCode": None,
            "options": None,
        }

    @override
    async def attach_output_queue(self, run_id: str | None) -> None:
        return

    @override
    def resume_output_queue(self) -> None:
        return

    @override
    def next_output_queue(self) -> None:
        return

    @override
    async def read_output(self) -> None:
        return


async def prepare_krunner_env(_local_config: Mapping[str, Any]) -> Mapping[str, str]:
    return {}
