from __future__ import annotations

import logging
import threading
from typing import TYPE_CHECKING, Any, ClassVar, Final
import pyarrow as pa
import pandas as pd
import clickhouse_connect
from pydantic import SecretStr
from contextlib import suppress

from sibi_flux.core import SecureResource
from sibi_flux.core.type_maps import DASK_TO_CLICKHOUSE_DTYPE
from sibi_flux.dask_cluster.core import safe_compute
from contextlib import nullcontext

try:
    from tqdm.dask import TqdmCallback
except ImportError:
    TqdmCallback = None

if TYPE_CHECKING:
    import dask.dataframe as dd

# Global thread-local storage for ClickHouse clients on workers
_worker_clients = threading.local()

def _to_bool(val: Any) -> bool:
    if isinstance(val, bool):
        return val
    if isinstance(val, (int, float)):
        return bool(val)
    if isinstance(val, str):
        return val.strip().lower() in ("1", "true", "yes", "on")
    return False


class ClickHouseWriter(SecureResource):
    """
    High-performance ClickHouse writer using Native Arrow Stream.

    Features:
      - Direct PyArrow insertion (Zero-Copy)
      - Automatic Schema Evolution (Nullable inference)
      - Thread-local connection pooling for workers
    """

    # Create a local copy to avoid mutating the global type map
    ARROW_TO_CLICKHOUSE: ClassVar[dict[str, str]] = {
        **DASK_TO_CLICKHOUSE_DTYPE,
        "double": "Float64",
        "float": "Float32",
    }

    def __init__(
        self,
        *,
        host: str = "localhost",
        port: int | str = 8123,
        database: str = "default",
        user: str | SecretStr = "default",
        password: str | SecretStr = "",
        secure: bool | str = False,
        verify: bool | str = False,
        table: str = "test_table",
        order_by: str = "tuple()",
        engine: str | None = None,
        max_workers: int | str = 4,
        overwrite: bool | str = False,
        connect_timeout: int | str = 30,
        send_receive_timeout: int | str = 300,
        dask_client: Any | None = None,
        show_progress: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.host = host
        self.port = int(port)
        self.database = database
        self.user = user if isinstance(user, SecretStr) else SecretStr(user)
        self.password = (
            password if isinstance(password, SecretStr) else SecretStr(password)
        )
        self.secure = _to_bool(secure)
        self.verify = _to_bool(verify)
        self.table = table
        self.order_by = order_by
        self.engine = engine
        self.max_workers = int(max_workers)
        self.overwrite = _to_bool(overwrite)
        self.connect_timeout = int(connect_timeout)
        self.send_receive_timeout = int(send_receive_timeout)
        self.dask_client = dask_client
        self.show_progress = show_progress

        # Local thr-local for the main process client
        self._tlocal = threading.local()

        if self.overwrite:
            self._command(f"DROP TABLE IF EXISTS {self._ident(self.table)}")

    def save_to_clickhouse(self, df: dd.DataFrame) -> None:
        import dask.dataframe as dd
        if not isinstance(df, dd.DataFrame):
            raise TypeError("Expected Dask DataFrame")

        df = df.reset_index(drop=True)
        meta_arrow_schema = pa.Schema.from_pandas(df._meta, preserve_index=False)

        self._ensure_table_exists(meta_arrow_schema)

        conn_kwargs: dict[str, Any] = {
            "host": self.host,
            "port": self.port,
            "database": self.database,
            "username": self.user.get_secret_value(),
            "password": self.password.get_secret_value(),
            "secure": self.secure,
            "verify": self.verify,
            "compress": True,
            "connect_timeout": self.connect_timeout,
            "send_receive_timeout": self.send_receive_timeout,
        }

        self.logger.info(f"Submitting distributed write for table '{self.table}'")

        try:
            results = df.map_partitions(
                self._write_partition_static,
                table_name=self.table,
                conn_kwargs=conn_kwargs,
                meta=("rows_written", "int64"),
            )

            ctx = TqdmCallback(desc=f"Saving to ClickHouse table '{self.table}'") if self.show_progress and TqdmCallback else nullcontext()
            with ctx:
                results_computed = safe_compute(results, dask_client=self.dask_client)
            
            total_rows = results_computed.sum()
            self.logger.info(f"Distributed write completed. Total rows: {total_rows}")

        except Exception as e:
            self.logger.error(f"Distributed write failed: {e}", exc_info=True)
            raise

    def _ensure_table_exists(self, schema: pa.Schema) -> None:
        col_defs = []

        for field in schema:
            ch_type = self._map_arrow_type(field.type)

            if field.nullable:
                ch_type = f"Nullable({ch_type})"

            col_defs.append(f"{self._ident(field.name)} {ch_type}")

        schema_sql = ", ".join(col_defs)

        if not self.engine:
            ordering = self.order_by
            if ordering and not ordering.startswith("(") and ordering != "tuple()":
                ordering = f"({ordering})"
            engine_sql = f"ENGINE = MergeTree ORDER BY {ordering} SETTINGS allow_nullable_key = 1"
        else:
            engine_sql = self.engine

        ddl = f"CREATE TABLE IF NOT EXISTS {self._ident(self.table)} ({schema_sql}) {engine_sql}"
        self._command(ddl)

    def _map_arrow_type(self, t: pa.DataType) -> str:
        if pa.types.is_date32(t):
            return "Date32"
        if pa.types.is_date64(t):
            return "DateTime64(3)"

        if isinstance(t, pa.TimestampType):
            if t.unit == "ns":
                return "DateTime64(9)"
            if t.unit == "us":
                return "DateTime64(6)"
            if t.unit == "ms":
                return "DateTime64(3)"
            return "DateTime"

        if pa.types.is_dictionary(t):
            return "String"

        t_str = str(t)
        return self.ARROW_TO_CLICKHOUSE.get(t_str, "String")

    @staticmethod
    def _write_partition_static(
        pdf: pd.DataFrame, table_name: str, conn_kwargs: dict[str, Any]
    ) -> int:
        if pdf.empty:
            return 0

        try:
            # Shared client across partition writes in the SAME thread (on workers)
            # This reduces connection overhead dramatically.
            client_key = f"{conn_kwargs['host']}_{conn_kwargs['port']}_{conn_kwargs['database']}"
            client = getattr(_worker_clients, client_key, None)
            
            if client is None:
                client = clickhouse_connect.get_client(**conn_kwargs)
                setattr(_worker_clients, client_key, client)

            arrow_table = pa.Table.from_pandas(pdf, nthreads=1, preserve_index=False)
            client.insert_arrow(table_name, arrow_table)

            return len(pdf)
        except Exception as e:
            # If a client stale, clear it
            if "client" in locals():
                with suppress(Exception):
                    setattr(_worker_clients, client_key, None)
            raise RuntimeError(f"Worker write failed for table {table_name}: {e}") from e

    def _get_client(self) -> Any:
        cli = getattr(self._tlocal, "client", None)
        if cli:
            return cli

        cli = clickhouse_connect.get_client(
            host=self.host,
            port=self.port,
            database=self.database,
            username=self.user.get_secret_value(),
            password=self.password.get_secret_value(),
            secure=self.secure,
            verify=self.verify,
            compress=True,
            connect_timeout=self.connect_timeout,
            send_receive_timeout=self.send_receive_timeout,
        )
        self._tlocal.client = cli
        return cli

    def _command(self, sql: str) -> None:
        self._get_client().command(sql)

    @staticmethod
    def _ident(name: str) -> str:
        return f"`{name}`" if not name.startswith("`") else name

    def _cleanup(self) -> None:
        cli = getattr(self._tlocal, "client", None)
        if cli:
            with suppress(Exception):
                cli.close()
            delattr(self._tlocal, "client")
