"""
aaiclick.data.ch_client - ClickHouse client protocol and factory dispatch.

Defines ChClient Protocol for type-safe usage across the codebase,
and lazy-imports the appropriate concrete client based on AAICLICK_CH_URL.
"""

from __future__ import annotations

import os
from collections.abc import Sequence
from contextvars import ContextVar
from pathlib import Path
from typing import Protocol, cast

from aaiclick.backend import is_chdb

from ..formats import open_export_writer
from ..models import QueryStats


class QueryResult(Protocol):
    """Protocol for ClickHouse query results."""

    result_rows: list[tuple]
    column_names: list[str]


class ChClient(Protocol):
    """Protocol for async ClickHouse client operations.

    Both ChdbClient and clickhouse-connect AsyncClient satisfy this protocol.

    ``parameters`` are ClickHouse ``{name:Type}`` placeholder bindings.
    Both backends support them natively: clickhouse-connect routes them
    through its own typed substitution, chdb forwards to ``session.query(
    sql, fmt, params=...)``. Prefer parameter binding over string
    interpolation for any value that might be large or user-derived.
    """

    async def command(
        self,
        query: str,
        settings: dict | None = None,
        parameters: dict | None = None,
    ) -> object: ...
    async def query(
        self,
        query: str,
        settings: dict | None = None,
        parameters: dict | None = None,
    ) -> QueryResult: ...
    async def insert(
        self,
        table: str,
        data: Sequence[Sequence],
        column_names: Sequence[str] | None = None,
        column_oriented: bool = False,
        column_type_names: Sequence[str] | None = None,
    ) -> None: ...
    async def close(self) -> None: ...


_ch_client_var: ContextVar[ChClient | None] = ContextVar("ch_client", default=None)

# VS Code's debug console evaluates each `await` in a fresh Python Context, so
# ContextVar-bound clients are invisible. Setting AAICLICK_DEBUGGER=1 enables a
# module-level chdb fallback (chdb only — the only backend that can be created
# synchronously without a running event loop).
_DEBUGGER_ENABLED = bool(os.environ.get("AAICLICK_DEBUGGER"))
_debug_ch_client: ChClient | None = None


def get_ch_client() -> ChClient:
    """Return the ClickHouse client for the active data context."""
    client = _ch_client_var.get()
    if client is None and _DEBUGGER_ENABLED and is_chdb():
        global _debug_ch_client
        if _debug_ch_client is None:
            from .chdb_client import create_chdb_client

            _debug_ch_client = create_chdb_client()
        client = _debug_ch_client
    if client is None:
        raise RuntimeError(
            "No active data or orch context — use 'async with data_context()' or 'async with orch_context()'"
        )
    return client


async def execute_for_stats(
    query: str,
    settings: dict | None = None,
    parameters: dict | None = None,
    client: ChClient | None = None,
) -> QueryStats:
    """Run *query* purely to capture its ClickHouse run statistics.

    Backs the discard terminal (:meth:`Object.execute` runs ``… FORMAT Null``)
    and materializing writes (``copy`` / operator ``INSERT … SELECT``) that want
    the row/byte counts ClickHouse reports. Backend-agnostic: both clients'
    ``command()`` return a summary for a body-less statement — a
    ``clickhouse-connect`` ``QuerySummary`` (HTTP) or a ``ChdbCommandSummary``
    (chdb) — which :meth:`QueryStats.from_clickhouse_summary` maps uniformly.

    ``client`` defaults to the active context's client; ingest / operator
    callers pass their explicit client.
    """
    client = client or get_ch_client()
    summary = await client.command(query, settings=settings, parameters=parameters)
    return QueryStats.from_clickhouse_summary(summary)


async def export_query_to_file(query: str, path: str, fmt: str) -> str:
    """Stream a ``SELECT`` query result to a local file in *fmt*.

    - **chdb (embedded):** ``INSERT INTO FUNCTION file('path', fmt) <query>``
      — the embedded engine writes directly to disk. Compression is inferred
      from the path suffix by ClickHouse's ``file()``.
    - **clickhouse-connect (remote HTTP):** ``raw_stream`` yields uncompressed
      formatted byte chunks via an async ``StreamContext``; they are written
      to the local file via :func:`aaiclick.data.formats.open_export_writer`,
      which re-applies ``.gz`` / ``.xz`` compression client-side so output is
      byte-equivalent across backends. Awaiting each chunk keeps the event
      loop responsive between network reads.

    ``INSERT INTO FUNCTION file()`` is unusable against a remote server
    because it would write to the server's ``user_files_path`` rather than
    the client's filesystem.
    """
    abs_path = str(Path(path).resolve())
    client = get_ch_client()
    if is_chdb():
        safe_path = abs_path.replace("'", "\\'")
        await client.command(f"INSERT INTO FUNCTION file('{safe_path}', '{fmt}') {query}")
        return abs_path

    async with await client.raw_stream(query=query, fmt=fmt) as stream:  # type: ignore[attr-defined]
        with open_export_writer(abs_path) as f:
            async for chunk in stream:
                f.write(chunk)
    return abs_path


async def create_ch_client() -> ChClient:
    """Create a ClickHouse client from AAICLICK_CH_URL."""
    if is_chdb():
        from .chdb_client import create_chdb_client

        return create_chdb_client()

    from .clickhouse_client import create_clickhouse_client

    return cast(ChClient, await create_clickhouse_client())
