"""ClickHouse sink using clickhouse-connect."""

from __future__ import annotations

import json
from typing import Any

from .base import BaseSink

try:
    import clickhouse_connect  # type: ignore[import-untyped]
except ImportError:
    clickhouse_connect = None  # type: ignore[assignment]

_DDL = """
CREATE TABLE IF NOT EXISTS {database}.{table} (
    id String,
    element_type LowCardinality(String),
    text String,
    source_filename String,
    source_sha256 String,
    source_media_type LowCardinality(String),
    page Nullable(Int32),
    sequence_index Int32,
    section_path Array(String),
    section_path_tier Int8,
    schema_version LowCardinality(String),
    extraction_method LowCardinality(String),
    confidence Nullable(Float32),
    token_count Nullable(Int32),
    entities String DEFAULT '[]',
    classification String DEFAULT '{{}}',
    relations String DEFAULT '[]',
    structured_fields String DEFAULT '{{}}',
    table_data String DEFAULT '{{}}',
    created_at DateTime DEFAULT now()
) ENGINE = MergeTree()
ORDER BY (source_sha256, sequence_index)
PARTITION BY source_media_type
SETTINGS index_granularity = 8192
"""

_COLUMNS = [
    "id",
    "element_type",
    "text",
    "source_filename",
    "source_sha256",
    "source_media_type",
    "page",
    "sequence_index",
    "section_path",
    "section_path_tier",
    "schema_version",
    "extraction_method",
    "confidence",
    "token_count",
    "entities",
    "classification",
    "relations",
    "structured_fields",
    "table_data",
]


def _flatten(record: dict[str, Any]) -> dict[str, Any]:
    return {
        "id": record.get("id", ""),
        "element_type": record.get("element_type", ""),
        "text": record.get("text", ""),
        "source_filename": record.get("source_filename", ""),
        "source_sha256": record.get("source_sha256", ""),
        "source_media_type": record.get("source_media_type", ""),
        "page": record.get("page"),
        "sequence_index": record.get("sequence_index", 0),
        "section_path": record.get("section_path") or [],
        "section_path_tier": record.get("section_path_tier", 0),
        "schema_version": record.get("schema_version", ""),
        "extraction_method": record.get("extraction_method", ""),
        "confidence": record.get("confidence"),
        "token_count": record.get("token_count"),
        "entities": json.dumps(record.get("entities") or []),
        "classification": json.dumps(record.get("classification") or {}),
        "relations": json.dumps(record.get("relations") or []),
        "structured_fields": json.dumps(record.get("structured_fields") or {}),
        "table_data": json.dumps(record.get("table_data") or {}),
    }


class ClickHouseSink(BaseSink):
    def __init__(
        self,
        host: str = "localhost",
        port: int = 8123,
        database: str = "extractor",
        table: str = "elements",
        batch_size: int = 1000,
        username: str = "default",
        password: str = "",
    ) -> None:
        if clickhouse_connect is None:
            raise ImportError("Run: pip install 'coresdk-extractor[clickhouse]'")
        self.host = host
        self.port = port
        self.database = database
        self.table = table
        self.batch_size = batch_size
        self.username = username
        self.password = password
        self._client: Any = None
        self._buffer: list[dict[str, Any]] = []

    def connect(self) -> None:
        self._client = clickhouse_connect.get_client(
            host=self.host,
            port=self.port,
            username=self.username,
            password=self.password,
        )
        self._client.command(f"CREATE DATABASE IF NOT EXISTS {self.database}")
        ddl = _DDL.format(database=self.database, table=self.table)
        self._client.command(ddl)

    def feed(self, record: dict[str, Any]) -> None:
        self._buffer.append(record)
        if len(self._buffer) >= self.batch_size:
            self.write_batch(self._buffer)
            self._buffer = []

    def write_batch(self, records: list[dict[str, Any]]) -> None:
        if not records or self._client is None:
            return
        flat = [_flatten(r) for r in records]
        rows = [[row[col] for col in _COLUMNS] for row in flat]
        self._client.insert(
            f"{self.database}.{self.table}",
            rows,
            column_names=_COLUMNS,
        )

    def flush(self) -> None:
        if self._buffer:
            self.write_batch(self._buffer)
            self._buffer = []

    def close(self) -> None:
        if self._client is not None:
            self._client.close()
            self._client = None
