from __future__ import annotations

import asyncio
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime

import clickhouse_connect
from clickhouse_connect.driver.asyncclient import AsyncClient


DDL_DATABASE = "CREATE DATABASE IF NOT EXISTS relay"

DDL_MESSAGES = """
CREATE TABLE IF NOT EXISTS relay.messages (
    message_id   String,
    group_id     LowCardinality(String),
    from_peer    String,
    from_display String,
    content      String,
    timestamp    DateTime64(6, 'UTC'),
    seq          UInt64
) ENGINE = MergeTree
ORDER BY (group_id, seq)
PARTITION BY group_id
"""

DDL_PEERS = """
CREATE TABLE IF NOT EXISTS relay.peers (
    peer_id           String,
    group_id          LowCardinality(String),
    hostname          String,
    working_directory String,
    git_branch        Nullable(String),
    deleted           UInt8 DEFAULT 0,
    updated_at        DateTime64(6, 'UTC') DEFAULT now64(6)
) ENGINE = ReplacingMergeTree(updated_at, deleted)
ORDER BY (group_id, peer_id)
"""

DDL_CURSORS = """
CREATE TABLE IF NOT EXISTS relay.cursors (
    peer_id    String,
    group_id   LowCardinality(String),
    last_seq   UInt64,
    updated_at DateTime64(6, 'UTC') DEFAULT now64(6)
) ENGINE = ReplacingMergeTree(updated_at)
ORDER BY (group_id, peer_id)
"""

DDL_GROUPS = """
CREATE TABLE IF NOT EXISTS relay.groups (
    group_id      String,
    password_hash String,
    created_at    DateTime64(6, 'UTC') DEFAULT now64(6),
    deleted       UInt8 DEFAULT 0,
    updated_at    DateTime64(6, 'UTC') DEFAULT now64(6)
) ENGINE = ReplacingMergeTree(updated_at, deleted)
ORDER BY group_id
"""


@dataclass(frozen=True)
class ClickHouseConfig:
    host: str = "localhost"
    port: int = 8123
    username: str = "default"
    password: str = ""
    database: str = "relay"
    secure: bool = False


@dataclass(frozen=True)
class PeerRow:
    peer_id: str
    group_id: str
    hostname: str
    working_directory: str
    git_branch: str | None


@dataclass(frozen=True)
class MessageRow:
    message_id: str
    group_id: str
    from_peer: str
    from_display: str
    content: str
    timestamp: datetime
    seq: int


class RelayDB:
    _client: AsyncClient

    def __init__(self, config: ClickHouseConfig) -> None:
        self._config = config
        self._seq_locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock)

    async def connect(self) -> None:
        self._client = await clickhouse_connect.get_async_client(
            host=self._config.host,
            port=self._config.port,
            username=self._config.username,
            password=self._config.password,
            secure=self._config.secure,
        )

    async def close(self) -> None:
        await self._client.close()

    async def ensure_schema(self) -> None:
        await self._client.command(DDL_DATABASE)
        await self._client.command(DDL_MESSAGES)
        await self._client.command(DDL_PEERS)
        await self._client.command(DDL_CURSORS)
        await self._client.command(DDL_GROUPS)

    async def upsert_peer(
        self,
        peer_id: str,
        group_id: str,
        hostname: str,
        working_directory: str,
        git_branch: str | None,
    ) -> None:
        await self._client.command(
            "INSERT INTO relay.peers (peer_id, group_id, hostname, working_directory, git_branch, deleted) "
            "VALUES ({p_id:String}, {g_id:String}, {host:String}, {wd:String}, {branch:Nullable(String)}, 0)",
            parameters={
                "p_id": peer_id,
                "g_id": group_id,
                "host": hostname,
                "wd": working_directory,
                "branch": git_branch,
            },
        )

    async def delete_peer(self, peer_id: str, group_id: str) -> None:
        await self._client.command(
            "INSERT INTO relay.peers (peer_id, group_id, hostname, working_directory, git_branch, deleted) "
            "VALUES ({p_id:String}, {g_id:String}, '', '', NULL, 1)",
            parameters={"p_id": peer_id, "g_id": group_id},
        )

    async def get_peers(self, group_id: str) -> list[PeerRow]:
        result = await self._client.query(
            "SELECT peer_id, group_id, hostname, working_directory, git_branch "
            "FROM relay.peers FINAL "
            "WHERE group_id = {g_id:String} AND deleted = 0",
            parameters={"g_id": group_id},
        )
        return [
            PeerRow(
                peer_id=row[0],
                group_id=row[1],
                hostname=row[2],
                working_directory=row[3],
                git_branch=row[4],
            )
            for row in result.result_rows
        ]

    async def insert_message(
        self,
        message_id: str,
        group_id: str,
        from_peer: str,
        from_display: str,
        content: str,
        timestamp: datetime,
    ) -> int:
        async with self._seq_locks[group_id]:
            await self._client.command(
                "INSERT INTO relay.messages "
                "(message_id, group_id, from_peer, from_display, content, timestamp, seq) "
                "SELECT "
                "{m_id:String}, {g_id:String}, {fp:String}, {fd:String}, "
                "{ct:String}, {ts:DateTime64(6, 'UTC')}, "
                "coalesce(max(seq), 0) + 1 "
                "FROM relay.messages "
                "WHERE group_id = {g_id:String}",
                parameters={
                    "m_id": message_id,
                    "g_id": group_id,
                    "fp": from_peer,
                    "fd": from_display,
                    "ct": content,
                    "ts": timestamp,
                },
            )
            result = await self._client.query(
                "SELECT max(seq) FROM relay.messages "
                "WHERE group_id = {g_id:String} AND message_id = {m_id:String}",
                parameters={"g_id": group_id, "m_id": message_id},
            )
            return int(result.result_rows[0][0])

    async def get_messages_after(
        self, group_id: str, after_seq: int
    ) -> list[MessageRow]:
        result = await self._client.query(
            "SELECT message_id, group_id, from_peer, from_display, "
            "content, timestamp, seq "
            "FROM relay.messages "
            "WHERE group_id = {g_id:String} AND seq > {aseq:UInt64} "
            "ORDER BY seq ASC",
            parameters={"g_id": group_id, "aseq": after_seq},
        )
        return [self._row_to_message(row) for row in result.result_rows]

    async def get_all_messages(self, group_id: str) -> list[MessageRow]:
        result = await self._client.query(
            "SELECT message_id, group_id, from_peer, from_display, "
            "content, timestamp, seq "
            "FROM relay.messages "
            "WHERE group_id = {g_id:String} "
            "ORDER BY seq ASC",
            parameters={"g_id": group_id},
        )
        return [self._row_to_message(row) for row in result.result_rows]

    async def get_cursor(self, peer_id: str, group_id: str) -> int:
        result = await self._client.query(
            "SELECT last_seq "
            "FROM relay.cursors FINAL "
            "WHERE peer_id = {p_id:String} AND group_id = {g_id:String}",
            parameters={"p_id": peer_id, "g_id": group_id},
        )
        if not result.result_rows:
            return 0
        return int(result.result_rows[0][0])

    async def set_cursor(
        self, peer_id: str, group_id: str, last_seq: int
    ) -> None:
        await self._client.command(
            "INSERT INTO relay.cursors (peer_id, group_id, last_seq) "
            "VALUES ({p_id:String}, {g_id:String}, {ls:UInt64})",
            parameters={
                "p_id": peer_id,
                "g_id": group_id,
                "ls": last_seq,
            },
        )

    async def create_group(self, group_id: str, password_hash: str) -> None:
        await self._client.command(
            "INSERT INTO relay.groups (group_id, password_hash) "
            "VALUES ({g_id:String}, {ph:String})",
            parameters={"g_id": group_id, "ph": password_hash},
        )

    async def list_groups(self) -> list[str]:
        result = await self._client.query(
            "SELECT group_id FROM relay.groups FINAL WHERE deleted = 0"
        )
        return [row[0] for row in result.result_rows]

    async def get_group_password_hash(self, group_id: str) -> str | None:
        result = await self._client.query(
            "SELECT password_hash FROM relay.groups FINAL "
            "WHERE group_id = {g_id:String} AND deleted = 0",
            parameters={"g_id": group_id},
        )
        if not result.result_rows:
            return None
        return str(result.result_rows[0][0])

    async def group_exists(self, group_id: str) -> bool:
        return (await self.get_group_password_hash(group_id)) is not None

    @staticmethod
    def _row_to_message(row: tuple) -> MessageRow:  # type: ignore[type-arg]
        return MessageRow(
            message_id=row[0],
            group_id=row[1],
            from_peer=row[2],
            from_display=row[3],
            content=row[4],
            timestamp=row[5],
            seq=int(row[6]),
        )
