client.py

from __future__ import annotations

from dataclasses import asdict, dataclass
import time
from typing import Any, Iterator

from urllib3.exceptions import HTTPError as Urllib3HTTPError
from urllib3.exceptions import ProtocolError

from hotdata import ApiClient, Configuration
from hotdata.api.connections_api import ConnectionsApi
from hotdata.api.information_schema_api import InformationSchemaApi
from hotdata.api.query_api import QueryApi
from hotdata.api.query_runs_api import QueryRunsApi
from hotdata.api.results_api import ResultsApi
from hotdata.api.uploads_api import UploadsApi
from hotdata.exceptions import ApiException
from hotdata.models.async_query_response import AsyncQueryResponse
from hotdata.models.query_request import QueryRequest
from hotdata.models.query_response import QueryResponse
from hotdata.models.load_managed_table_request import LoadManagedTableRequest
from hotdata.models.table_info import TableInfo

from hotdata_runtime.env import (
    default_api_key,
    default_host,
    default_session_id,
    normalize_host,
    pick_workspace,
)
from hotdata_runtime.databases import (
    DEFAULT_SCHEMA,
    LoadManagedTableResult,
    ManagedDatabase,
    ManagedTable,
    MANAGED_SOURCE_TYPE,
    api_error_message,
    create_connection_request,
    is_parquet_path,
    managed_database_from_connection,
)
from hotdata_runtime.http import default_http_retries
from hotdata_runtime.result import QueryResult

_TERMINAL = frozenset({"succeeded", "failed", "cancelled"})
_RESULT_FAILURE = frozenset({"failed", "cancelled"})


@dataclass(frozen=True)
class ResultSummary:
    result_id: str
    status: str
    created_at: str | None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class RunHistoryItem:
    query_run_id: str
    status: str
    created_at: str | None
    execution_time_ms: int | None
    result_id: str | None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


class HotdataClient:
    """Thin wrapper around the Hotdata Python SDK with query polling helpers."""

    def __init__(
        self,
        api_key: str,
        workspace_id: str,
        *,
        host: str | None = None,
        session_id: str | None = None,
    ) -> None:
        self._host = normalize_host(host) if host else default_host()
        self._api_key = api_key
        self._workspace_id = workspace_id
        self._session_id = session_id
        self._config = Configuration(
            host=self._host,
            api_key=api_key,
            workspace_id=workspace_id,
            session_id=session_id,
            retries=default_http_retries(),
        )
        self._api = ApiClient(self._config)

    @classmethod
    def from_env(cls) -> HotdataClient:
        api_key = default_api_key()
        if not api_key:
            raise RuntimeError("HOTDATA_API_KEY must be set.")
        host = default_host()
        session = default_session_id()
        workspace_id = pick_workspace(api_key, host, session)
        return cls(api_key, workspace_id, host=host, session_id=session)

    @property
    def workspace_id(self) -> str:
        return self._workspace_id

    @property
    def host(self) -> str:
        return self._host

    @property
    def session_id(self) -> str | None:
        return self._session_id

    @property
    def api(self) -> ApiClient:
        return self._api

    def close(self) -> None:
        self._api.close()

    def __enter__(self) -> HotdataClient:
        return self

    def __exit__(self, *args: object) -> None:
        self.close()

    def connections(self) -> ConnectionsApi:
        return ConnectionsApi(self._api)

    def _information_schema(self) -> InformationSchemaApi:
        return InformationSchemaApi(self._api)

    def _query_api(self) -> QueryApi:
        return QueryApi(self._api)

    def _query_runs_api(self) -> QueryRunsApi:
        return QueryRunsApi(self._api)

    def _results_api(self) -> ResultsApi:
        return ResultsApi(self._api)

    def query_runs(self) -> QueryRunsApi:
        return self._query_runs_api()

    def results(self) -> ResultsApi:
        return self._results_api()

    def uploads(self) -> UploadsApi:
        return UploadsApi(self._api)

    def list_managed_databases(self) -> list[ManagedDatabase]:
        listing = self.connections().list_connections()
        return [
            managed_database_from_connection(c)
            for c in listing.connections
            if c.source_type == MANAGED_SOURCE_TYPE
        ]

    def resolve_managed_database(self, name_or_id: str) -> ManagedDatabase:
        listing = self.connections().list_connections()
        match = None
        for c in listing.connections:
            if c.id == name_or_id or c.name == name_or_id:
                match = c
                break
        if match is None:
            raise KeyError(f"No database named or with id {name_or_id!r}")
        if match.source_type != MANAGED_SOURCE_TYPE:
            raise ValueError(
                f"{match.name!r} is not a managed database "
                f"(source_type: {match.source_type})"
            )
        return managed_database_from_connection(match)

    def create_managed_database(
        self,
        name: str,
        *,
        schema: str = DEFAULT_SCHEMA,
        tables: list[str] | None = None,
    ) -> ManagedDatabase:
        request = create_connection_request(name, schema=schema, tables=tables)
        try:
            created = self.connections().create_connection(request)
        except ApiException as e:
            raise RuntimeError(api_error_message(e)) from e
        return managed_database_from_connection(created)

    def delete_managed_database(self, name_or_id: str) -> None:
        db = self.resolve_managed_database(name_or_id)
        try:
            self.connections().delete_connection(db.id)
        except ApiException as e:
            raise RuntimeError(api_error_message(e)) from e

    def list_managed_tables(
        self,
        database: str,
        *,
        schema: str | None = None,
    ) -> list[ManagedTable]:
        db = self.resolve_managed_database(database)
        rows: list[ManagedTable] = []
        for t in self.iter_tables(connection_id=db.id):
            if schema is not None and t.var_schema != schema:
                continue
            rows.append(
                ManagedTable(
                    full_name=f"{db.name}.{t.var_schema}.{t.table}",
                    schema=t.var_schema,
                    table=t.table,
                    synced=t.synced,
                    last_sync=t.last_sync,
                )
            )
        rows.sort(key=lambda row: (row.schema, row.table))
        return rows

    def upload_parquet(self, path: str) -> str:
        if not is_parquet_path(path):
            raise ValueError(
                f"Managed table loads require a parquet file (got {path!r})"
            )
        with open(path, "rb") as f:
            data = f.read()
        try:
            uploaded = self.uploads().upload_file(
                data,
                _content_type="application/octet-stream",
            )
        except ApiException as e:
            raise RuntimeError(api_error_message(e)) from e
        return uploaded.id

    def load_managed_table(
        self,
        database: str,
        table: str,
        *,
        schema: str = DEFAULT_SCHEMA,
        upload_id: str | None = None,
        file: str | None = None,
    ) -> LoadManagedTableResult:
        if (upload_id is None) == (file is None):
            raise ValueError("Exactly one of upload_id or file is required")
        db = self.resolve_managed_database(database)
        if upload_id is not None:
            resolved_upload_id = upload_id
        else:
            assert file is not None
            resolved_upload_id = self.upload_parquet(file)
        request = LoadManagedTableRequest(
            mode="replace",
            upload_id=resolved_upload_id,
        )
        try:
            loaded = self.connections().load_managed_table(
                db.id,
                schema,
                table,
                request,
            )
        except ApiException as e:
            raise RuntimeError(api_error_message(e)) from e
        return LoadManagedTableResult(
            connection_id=loaded.connection_id,
            schema_name=loaded.schema_name,
            table_name=loaded.table_name,
            row_count=loaded.row_count,
            full_name=f"{db.name}.{loaded.schema_name}.{loaded.table_name}",
        )

    def delete_managed_table(
        self,
        database: str,
        table: str,
        *,
        schema: str = DEFAULT_SCHEMA,
    ) -> None:
        db = self.resolve_managed_database(database)
        try:
            self.connections().delete_managed_table(db.id, schema, table)
        except ApiException as e:
            raise RuntimeError(api_error_message(e)) from e

    def list_recent_results(
        self,
        *,
        limit: int = 50,
        offset: int = 0,
    ) -> list[ResultSummary]:
        listing = self.results().list_results(limit=limit, offset=offset)
        return [
            ResultSummary(
                result_id=r.id,
                status=r.status,
                created_at=r.created_at,
            )
            for r in listing.results
        ]

    def list_run_history(
        self,
        *,
        limit: int = 20,
    ) -> list[RunHistoryItem]:
        listing = self.query_runs().list_query_runs(limit=limit)
        return [
            RunHistoryItem(
                query_run_id=r.id,
                status=r.status,
                created_at=r.created_at,
                execution_time_ms=r.execution_time_ms,
                result_id=r.result_id,
            )
            for r in listing.query_runs
        ]

    def iter_tables(
        self,
        *,
        connection_id: str | None = None,
        include_columns: bool = False,
        page_size: int = 200,
    ) -> Iterator[TableInfo]:
        cursor: str | None = None
        while True:
            resp = self._information_schema().information_schema(
                connection_id=connection_id,
                include_columns=include_columns,
                limit=page_size,
                cursor=cursor,
            )
            yield from resp.tables
            if not resp.has_more or not resp.next_cursor:
                break
            cursor = resp.next_cursor

    def qualified_table_name(self, t: TableInfo) -> str:
        return f"{t.connection}.{t.var_schema}.{t.table}"

    def list_qualified_table_names(
        self, *, limit: int = 5000, connection_id: str | None = None
    ) -> list[str]:
        out: list[str] = []
        for t in self.iter_tables(connection_id=connection_id):
            out.append(self.qualified_table_name(t))
            if len(out) >= limit:
                break
        return sorted(out)

    def connection_id_by_name(self) -> dict[str, str]:
        listing = self.connections().list_connections()
        id_map: dict[str, str] = {}
        duplicate_names: set[str] = set()
        for c in listing.connections:
            if c.name in id_map and id_map[c.name] != c.id:
                duplicate_names.add(c.name)
            id_map[c.name] = c.id
        if duplicate_names:
            names = ", ".join(sorted(duplicate_names))
            raise RuntimeError(
                f"Duplicate connection names found: {names}. "
                "Use an explicit connection_id."
            )
        return id_map

    def columns_for_qualified(
        self,
        qualified: str,
        *,
        connection_id: str | None = None,
    ) -> list[TableInfo]:
        parts = qualified.split(".")
        if len(parts) < 3:
            raise ValueError(
                f"Expected connection.schema.table, got {qualified!r}"
            )
        conn_name, schema_name, table_name = (
            parts[0],
            parts[1],
            ".".join(parts[2:]),
        )
        conn_id = connection_id
        if conn_id is None:
            id_map = self.connection_id_by_name()
            conn_id = id_map.get(conn_name)
            if not conn_id:
                raise KeyError(f"Unknown connection {conn_name!r}")
        resp = self._information_schema().information_schema(
            connection_id=conn_id,
            var_schema=schema_name,
            table=table_name,
            include_columns=True,
            limit=10,
        )
        if not resp.tables:
            return []
        first = resp.tables[0]
        return first.columns or []

    def _poll_query_run(
        self,
        query_run_id: str,
        *,
        timeout_s: float = 300.0,
        interval_s: float = 0.5,
    ):
        runs = self._query_runs_api()
        deadline = time.monotonic() + timeout_s
        last = None
        while time.monotonic() < deadline:
            last = runs.get_query_run(query_run_id)
            if last.status in _TERMINAL:
                return last
            time.sleep(interval_s)
        raise TimeoutError(
            f"Query run {query_run_id} did not finish within {timeout_s}s "
            f"(last status: {getattr(last, 'status', None)})"
        )

    def _wait_result_ready(
        self,
        result_id: str,
        *,
        timeout_s: float = 300.0,
        interval_s: float = 0.5,
    ):
        results = self._results_api()
        deadline = time.monotonic() + timeout_s
        last = None
        while time.monotonic() < deadline:
            last = results.get_result(result_id)
            if last.status == "ready":
                return last
            if last.status in _RESULT_FAILURE:
                raise RuntimeError(
                    last.error_message or f"Result {last.status}"
                )
            time.sleep(interval_s)
        raise TimeoutError(
            f"Result {result_id} not ready within {timeout_s}s "
            f"(last status: {getattr(last, 'status', None)})"
        )

    def execute_sql(self, sql: str) -> QueryResult:
        last_err: BaseException | None = None
        for attempt in range(3):
            try:
                return self._execute_sql_once(sql)
            except (ProtocolError, ConnectionResetError, Urllib3HTTPError) as e:
                last_err = e
                if attempt == 2:
                    raise
                time.sleep(0.2 * (2**attempt))
        raise last_err  # pragma: no cover

    def _execute_sql_once(self, sql: str) -> QueryResult:
        q = self._query_api()
        try:
            raw = q.query(QueryRequest(sql=sql))
        except ApiException as e:
            raise RuntimeError(e.reason or str(e)) from e

        if isinstance(raw, AsyncQueryResponse):
            run = self._poll_query_run(raw.query_run_id)
            if run.status != "succeeded":
                raise RuntimeError(
                    run.error_message or f"Query failed ({run.status})"
                )
            if run.result_id:
                persisted = self._wait_result_ready(run.result_id)
                return QueryResult.from_get_result(persisted)
            raise RuntimeError("Query succeeded but no result_id was returned.")

        if isinstance(raw, QueryResponse):
            return QueryResult.from_query_response(raw)

        raise RuntimeError(f"Unexpected query response type: {type(raw)!r}")

    def get_result(self, result_id: str) -> QueryResult:
        r = self._results_api().get_result(result_id)
        if r.status != "ready":
            r = self._wait_result_ready(result_id)
        return QueryResult.from_get_result(r)


def from_env() -> HotdataClient:
    return HotdataClient.from_env()