from __future__ import annotations

import asyncio
import json
import logging
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import TypeVar

import websockets.client
import websockets.exceptions
from exponent.commands.utils import ConnectionTracker
from exponent.core.remote_execution import files, git, system_context
from exponent.core.remote_execution.code_execution import execute_code
from exponent.core.remote_execution.command_execution import execute_command
from exponent.core.remote_execution.file_write import execute_file_write
from exponent.core.remote_execution.session import (
    RemoteExecutionClientSession,
    get_session,
)
from exponent.core.remote_execution.types import (
    CLIConnectedState,
    CodeExecutionRequest,
    CommandRequest,
    CreateChatResponse,
    ExecutionEndResponse,
    FileWriteRequest,
    GetAllTrackedFilesRequest,
    GetFileAttachmentRequest,
    GetFileAttachmentsRequest,
    GetMatchingFilesRequest,
    HeartbeatInfo,
    ListFilesRequest,
    RemoteExecutionRequestType,
    RemoteExecutionResponseType,
    StartChatRequest,
    StartChatResponse,
    SystemContextRequest,
    UseToolsConfig,
)
from exponent.core.remote_execution.utils import (
    assert_unreachable,
    deserialize_api_response,
    deserialize_request_data,
    serialize_message,
)
from httpx import (
    AsyncClient,
    codes as http_status,
)
from pydantic import BaseModel

logger = logging.getLogger(__name__)


TModel = TypeVar("TModel", bound=BaseModel)


class RemoteExecutionClient:
    def __init__(
        self,
        session: RemoteExecutionClientSession,
    ):
        self.current_session = session

        self.file_cache: files.FileCache = files.FileCache(session.working_directory)

    @property
    def working_directory(self) -> str:
        return self.current_session.working_directory

    @property
    def api_client(self) -> AsyncClient:
        return self.current_session.api_client

    async def run_connection(  # noqa: PLR0915,PLR0912
        self,
        chat_uuid: str,
        connection_tracker: ConnectionTracker | None = None,
    ) -> None:
        stop = False

        async for websocket in self.ws_connect(
            f"/api/ws/chat/{chat_uuid}",
        ):
            if connection_tracker is not None:
                await connection_tracker.set_connected(True)

            beats: asyncio.Queue[HeartbeatInfo] = asyncio.Queue()
            requests: asyncio.Queue[RemoteExecutionRequestType] = asyncio.Queue()
            results: asyncio.Queue[RemoteExecutionResponseType] = asyncio.Queue()

            async def beat() -> None:
                while True:
                    info = self.get_heartbeat_info()
                    await beats.put(info)
                    await asyncio.sleep(3)

            async def executor() -> None:
                while True:
                    request = await requests.get()
                    response = await self.handle_request(request)
                    await results.put(response)

            beat_task = asyncio.create_task(beat())
            executor_task = asyncio.create_task(executor())

            try:
                while True:
                    recv = asyncio.create_task(websocket.recv())
                    get_beat = asyncio.create_task(beats.get())
                    get_result = asyncio.create_task(results.get())

                    done, pending = await asyncio.wait(
                        [recv, get_beat, get_result],
                        return_when=asyncio.FIRST_COMPLETED,
                    )

                    for task in pending:
                        task.cancel()

                        try:
                            await task
                        except asyncio.CancelledError:
                            pass

                    if recv in done:
                        msg = json.loads(str(recv.result()))

                        if msg["type"] == "request":
                            data = json.dumps(msg["data"])
                            request = deserialize_request_data(data)
                            await requests.put(request)

                    elif get_beat in done:
                        info = get_beat.result()
                        data = json.loads(info.model_dump_json())
                        msg = json.dumps({"type": "heartbeat", "data": data})
                        await websocket.send(msg)

                    elif get_result in done:
                        response = get_result.result()
                        data = json.loads(serialize_message(response))
                        msg = json.dumps({"type": "result", "data": data})
                        await websocket.send(msg)

            except websockets.exceptions.ConnectionClosed as e:
                if e.rcvd is not None and e.rcvd.code == 1000:  # noqa: PLR2004
                    stop = True
            except TimeoutError:
                pass

            beat_task.cancel()
            executor_task.cancel()

            if stop:
                break

            if connection_tracker is not None:
                await connection_tracker.set_connected(False)

    async def check_remote_end_event(self, chat_uuid: str) -> bool:
        response = await self.api_client.get(
            f"/api/remote_execution/{chat_uuid}/execution_end",
        )
        execution_end_response = await deserialize_api_response(
            response, ExecutionEndResponse
        )
        return execution_end_response.execution_ended

    async def create_chat(self) -> CreateChatResponse:
        response = await self.api_client.post(
            "/api/remote_execution/create_chat",
        )
        return await deserialize_api_response(response, CreateChatResponse)

    async def start_chat(
        self, chat_uuid: str, prompt: str, use_tools_config: UseToolsConfig
    ) -> StartChatResponse:
        response = await self.api_client.post(
            "/api/remote_execution/start_chat",
            json=StartChatRequest(
                chat_uuid=chat_uuid,
                prompt=prompt,
                use_tools_config=use_tools_config,
            ).model_dump(),
            timeout=60,
        )
        return await deserialize_api_response(response, StartChatResponse)

    def get_heartbeat_info(self) -> HeartbeatInfo:
        return HeartbeatInfo(
            system_info=system_context.get_system_info(self.working_directory),
        )

    async def send_heartbeat(self, chat_uuid: str) -> CLIConnectedState:
        logger.info(f"Sending heartbeat for chat_uuid {chat_uuid}")
        heartbeat_info = self.get_heartbeat_info()
        response = await self.api_client.post(
            f"/api/remote_execution/{chat_uuid}/heartbeat",
            content=heartbeat_info.model_dump_json(),
            timeout=60,
        )
        if response.status_code != http_status.OK:
            raise Exception(
                f"Heartbeat failed with status code {response.status_code} and response {response.text}"
            )
        connected_state = await deserialize_api_response(response, CLIConnectedState)
        logger.info(f"Heartbeat response: {connected_state}")
        return connected_state

    async def handle_request(
        self, request: RemoteExecutionRequestType
    ) -> RemoteExecutionResponseType:
        response: RemoteExecutionResponseType
        if isinstance(request, CodeExecutionRequest):
            response = await execute_code(
                request, self.current_session, working_directory=self.working_directory
            )
        elif isinstance(request, FileWriteRequest):
            response = execute_file_write(
                request, working_directory=self.working_directory
            )
        elif isinstance(request, ListFilesRequest):
            response = files.list_files(request)
        elif isinstance(request, GetFileAttachmentRequest):
            response = files.get_file_attachment(request)
        elif isinstance(request, GetFileAttachmentsRequest):
            response = files.get_file_attachments(request)
        elif isinstance(request, GetMatchingFilesRequest):
            response = await files.get_matching_files(request, self.file_cache)
        elif isinstance(request, SystemContextRequest):
            response = system_context.get_system_context(
                request, self.working_directory
            )
        elif isinstance(request, GetAllTrackedFilesRequest):
            response = await git.get_all_tracked_files(request, self.working_directory)
        elif isinstance(request, CommandRequest):
            response = await execute_command(request, self.working_directory)
        else:
            assert_unreachable(request)
        return response

    def ws_connect(self, path: str) -> websockets.client.connect:
        base_url = (
            str(self.api_client.base_url)
            .replace("http://", "ws://")
            .replace("https://", "wss://")
        )

        url = f"{base_url}{path}"
        headers = {"api-key": self.api_client.headers["api-key"]}

        return websockets.client.connect(
            url, extra_headers=headers, timeout=10, ping_timeout=10
        )

    @staticmethod
    @asynccontextmanager
    async def session(
        api_key: str, base_url: str, working_directory: str
    ) -> AsyncGenerator[RemoteExecutionClient, None]:
        async with get_session(working_directory, base_url, api_key) as session:
            yield RemoteExecutionClient(session)
