"""SAP HANA Cloud checkpoint saver for LangGraph.

This module implements the BaseCheckpointSaver interface for SAP HANA Cloud,
allowing LangGraph agents to persist state, enabling conversation memory,
human-in-the-loop, time travel, and fault recovery on SAP BTP.

Typical usage:

    from langgraph_checkpoint_hana import HANASaver

    with HANASaver.from_conn_info(
        address="your-hana-host.hanacloud.ondemand.com",
        port=443,
        user="DBADMIN",
        password="...",
    ) as checkpointer:
        graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)
        result = graph.invoke({"messages": [("human", "hello")]}, config)
"""

from __future__ import annotations

import base64
import json
import logging
from contextlib import contextmanager
from typing import Any, AsyncIterator, Iterator, Optional, Sequence

from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
    WRITES_IDX_MAP,
    BaseCheckpointSaver,
    ChannelVersions,
    Checkpoint,
    CheckpointMetadata,
    CheckpointTuple,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer

logger = logging.getLogger(__name__)

# ──────────────────────────────────────────────────────────────────────
# DDL — table definitions for HANA Cloud
# ──────────────────────────────────────────────────────────────────────

_CHECKPOINTS_TABLE = "LANGGRAPH_CHECKPOINTS"
_WRITES_TABLE = "LANGGRAPH_CHECKPOINT_WRITES"

_CREATE_CHECKPOINTS = f"""
CREATE TABLE "{_CHECKPOINTS_TABLE}" (
    thread_id       NVARCHAR(256)  NOT NULL,
    checkpoint_ns   NVARCHAR(256)  NOT NULL DEFAULT '',
    checkpoint_id   NVARCHAR(256)  NOT NULL,
    parent_checkpoint_id NVARCHAR(256),
    type            NVARCHAR(64),
    checkpoint_data NCLOB          NOT NULL,
    metadata_data   NCLOB,
    created_at      TIMESTAMP      DEFAULT CURRENT_TIMESTAMP,
    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
)
"""

_CREATE_WRITES = f"""
CREATE TABLE "{_WRITES_TABLE}" (
    thread_id       NVARCHAR(256)  NOT NULL,
    checkpoint_ns   NVARCHAR(256)  NOT NULL DEFAULT '',
    checkpoint_id   NVARCHAR(256)  NOT NULL,
    task_id         NVARCHAR(256)  NOT NULL,
    task_path       NVARCHAR(512)  NOT NULL DEFAULT '',
    idx             INTEGER        NOT NULL,
    channel         NVARCHAR(256)  NOT NULL,
    type            NVARCHAR(64),
    blob_data       NCLOB,
    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
)
"""

# ──────────────────────────────────────────────────────────────────────
# SQL templates
# ──────────────────────────────────────────────────────────────────────

_SELECT_CHECKPOINT = f"""
    SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint_data,
           metadata_data
    FROM "{_CHECKPOINTS_TABLE}"
    WHERE thread_id = ? AND checkpoint_ns = ?
"""

_UPSERT_CHECKPOINT = f"""
    UPSERT "{_CHECKPOINTS_TABLE}"
    (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id,
     type, checkpoint_data, metadata_data, created_at)
    VALUES (?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
    WITH PRIMARY KEY
"""

_UPSERT_WRITE = f"""
    UPSERT "{_WRITES_TABLE}"
    (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path,
     idx, channel, type, blob_data)
    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
    WITH PRIMARY KEY
"""


class HANASaver(BaseCheckpointSaver):
    """LangGraph checkpoint saver backed by SAP HANA Cloud.

    Implements the full ``BaseCheckpointSaver`` interface so that any
    LangGraph graph compiled with this checkpointer will persist its
    state in HANA Cloud tables.

    Supports both synchronous and asynchronous usage (async delegates
    to sync because the ``hdbcli`` driver is synchronous).

    Example::

        with HANASaver.from_conn_info(
            address="host.hanacloud.ondemand.com",
            port=443,
            user="DBADMIN",
            password="secret",
        ) as saver:
            graph = workflow.compile(checkpointer=saver)
            graph.invoke(inputs, {"configurable": {"thread_id": "t1"}})
    """

    conn: Any  # hdbcli.dbapi.Connection

    def __init__(self, conn: Any) -> None:
        super().__init__(serde=JsonPlusSerializer())
        self.conn = conn

    # ── Factory / lifecycle ──────────────────────────────────────────

    @classmethod
    def from_conn_info(
        cls,
        *,
        address: str,
        port: int = 443,
        user: str,
        password: str,
        encrypt: bool = True,
        autocommit: bool = False,
        **hdbcli_kwargs: Any,
    ) -> "HANASaver":
        """Create a ``HANASaver`` from HANA connection parameters.

        Args:
            address: HANA Cloud host name.
            port: Port number (443 for HANA Cloud).
            user: Database user.
            password: Database password.
            encrypt: Use TLS (default ``True``; required for HANA Cloud).
            autocommit: HANA autocommit flag.
            **hdbcli_kwargs: Additional keyword arguments passed to
                ``hdbcli.dbapi.connect``.

        Returns:
            A new ``HANASaver`` instance with an open connection.
        """
        from hdbcli import dbapi

        conn = dbapi.connect(
            address=address,
            port=port,
            user=user,
            password=password,
            encrypt=encrypt,
            autocommit=autocommit,
            **hdbcli_kwargs,
        )
        return cls(conn=conn)

    @classmethod
    def from_env(cls) -> "HANASaver":
        """Create a ``HANASaver`` using environment variables.

        Reads ``HANA_HOST``, ``HANA_PORT``, ``HANA_USER``, ``HANA_PASSWORD``
        from the environment.  Convenient for containerised deployments on
        SAP BTP / Kyma.
        """
        import os

        return cls.from_conn_info(
            address=os.environ["HANA_HOST"],
            port=int(os.environ.get("HANA_PORT", "443")),
            user=os.environ["HANA_USER"],
            password=os.environ["HANA_PASSWORD"],
        )

    def __enter__(self) -> "HANASaver":
        self.setup()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        try:
            self.conn.close()
        except Exception:
            pass

    # ── Setup ────────────────────────────────────────────────────────

    def setup(self) -> None:
        """Create checkpoint tables if they do not exist.

        Safe to call multiple times — checks ``SYS.TABLES`` before
        issuing DDL.
        """
        cur = self.conn.cursor()
        try:
            for table_name, ddl in [
                (_CHECKPOINTS_TABLE, _CREATE_CHECKPOINTS),
                (_WRITES_TABLE, _CREATE_WRITES),
            ]:
                if not self._table_exists(cur, table_name):
                    cur.execute(ddl)
                    self.conn.commit()
                    logger.info("Created table %s", table_name)
                else:
                    logger.debug("Table %s already exists", table_name)
        finally:
            cur.close()

    @staticmethod
    def _table_exists(cur: Any, table_name: str) -> bool:
        """Check table existence via HANA system view."""
        cur.execute(
            'SELECT COUNT(*) FROM "SYS"."TABLES" WHERE TABLE_NAME = ?',
            (table_name,),
        )
        return cur.fetchone()[0] > 0

    # ── Helpers ──────────────────────────────────────────────────────

    @staticmethod
    def _cfg(config: RunnableConfig) -> tuple[str, str]:
        """Extract (thread_id, checkpoint_ns) from a runnable config."""
        c = config.get("configurable", {})
        return c.get("thread_id", ""), c.get("checkpoint_ns", "")

    def _dump(self, obj: Any) -> tuple[str, str]:
        """Serialise an object via serde → (type, base64_blob).

        JsonPlusSerializer produces msgpack bytes which are not valid
        UTF-8, so we base64-encode them for storage in NCLOB columns.
        """
        type_str, data_bytes = self.serde.dumps_typed(obj)
        if isinstance(data_bytes, bytes):
            return type_str, base64.b64encode(data_bytes).decode("ascii")
        return type_str, data_bytes

    def _load(self, type_str: str, blob: str) -> Any:
        """Deserialise a (type, base64_blob) pair back into an object."""
        if isinstance(blob, str):
            data_bytes = base64.b64decode(blob)
        else:
            data_bytes = blob
        return self.serde.loads_typed((type_str, data_bytes))

    # ── Core interface: get_tuple ────────────────────────────────────

    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
        """Fetch the checkpoint tuple for a given configuration.

        If ``checkpoint_id`` is present in the config, that exact
        checkpoint is returned.  Otherwise the most recent checkpoint
        for the thread is returned.
        """
        thread_id, ns = self._cfg(config)
        cp_id = config.get("configurable", {}).get("checkpoint_id")

        cur = self.conn.cursor()
        try:
            if cp_id:
                cur.execute(
                    _SELECT_CHECKPOINT + " AND checkpoint_id = ?",
                    (thread_id, ns, cp_id),
                )
            else:
                cur.execute(
                    _SELECT_CHECKPOINT
                    + " ORDER BY checkpoint_id DESC LIMIT 1",
                    (thread_id, ns),
                )

            row = cur.fetchone()
            if row is None:
                return None

            cp_id, parent_id, type_str, cp_blob, meta_blob = row

            checkpoint = self._load(type_str, cp_blob)
            metadata = json.loads(meta_blob) if meta_blob else {}
            pending = self._load_pending_writes(cur, thread_id, ns, cp_id)

            cfg = {
                "configurable": {
                    "thread_id": thread_id,
                    "checkpoint_ns": ns,
                    "checkpoint_id": cp_id,
                }
            }
            parent_cfg = (
                {
                    "configurable": {
                        "thread_id": thread_id,
                        "checkpoint_ns": ns,
                        "checkpoint_id": parent_id,
                    }
                }
                if parent_id
                else None
            )

            return CheckpointTuple(
                config=cfg,
                checkpoint=checkpoint,
                metadata=metadata,
                parent_config=parent_cfg,
                pending_writes=pending,
            )
        finally:
            cur.close()

    def _load_pending_writes(
        self,
        cur: Any,
        thread_id: str,
        ns: str,
        checkpoint_id: str,
    ) -> list[tuple[str, str, Any]]:
        """Load pending writes associated with a checkpoint."""
        cur.execute(
            f"""
            SELECT task_id, channel, type, blob_data
            FROM "{_WRITES_TABLE}"
            WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?
            ORDER BY idx
            """,
            (thread_id, ns, checkpoint_id),
        )
        result: list[tuple[str, str, Any]] = []
        for task_id, channel, type_str, blob in cur.fetchall():
            value = self._load(type_str, blob) if blob and type_str else None
            result.append((task_id, channel, value))
        return result

    # ── Core interface: list ─────────────────────────────────────────

    def list(
        self,
        config: Optional[RunnableConfig],
        *,
        filter: Optional[dict[str, Any]] = None,
        before: Optional[RunnableConfig] = None,
        limit: Optional[int] = None,
    ) -> Iterator[CheckpointTuple]:
        """List checkpoints matching the given criteria.

        Results are ordered newest-first (descending ``checkpoint_id``).

        Args:
            config: Filter by thread_id / checkpoint_ns.
            filter: Filter by metadata key-value pairs.
            before: Only return checkpoints older than this config's
                ``checkpoint_id``.
            limit: Maximum number of results.
        """
        thread_id, ns = self._cfg(config) if config else ("", "")

        cur = self.conn.cursor()
        try:
            where_clauses: list[str] = []
            params: list[Any] = []

            if thread_id:
                where_clauses.append("thread_id = ?")
                params.append(thread_id)
            if ns:
                where_clauses.append("checkpoint_ns = ?")
                params.append(ns)
            if before:
                before_id = before.get("configurable", {}).get("checkpoint_id")
                if before_id:
                    where_clauses.append("checkpoint_id < ?")
                    params.append(before_id)

            where_sql = (
                "WHERE " + " AND ".join(where_clauses)
                if where_clauses
                else ""
            )

            query = f"""
                SELECT checkpoint_id, parent_checkpoint_id, type,
                       checkpoint_data, metadata_data
                FROM "{_CHECKPOINTS_TABLE}"
                {where_sql}
                ORDER BY checkpoint_id DESC
            """
            if limit:
                query += f" LIMIT {int(limit)}"

            cur.execute(query, params)

            for cp_id, parent_id, type_str, cp_blob, meta_blob in cur.fetchall():
                checkpoint = self._load(type_str, cp_blob)
                metadata = json.loads(meta_blob) if meta_blob else {}

                # Apply metadata filter if provided
                if filter and not all(
                    metadata.get(k) == v for k, v in filter.items()
                ):
                    continue

                cfg = {
                    "configurable": {
                        "thread_id": thread_id,
                        "checkpoint_ns": ns,
                        "checkpoint_id": cp_id,
                    }
                }
                parent_cfg = (
                    {
                        "configurable": {
                            "thread_id": thread_id,
                            "checkpoint_ns": ns,
                            "checkpoint_id": parent_id,
                        }
                    }
                    if parent_id
                    else None
                )

                yield CheckpointTuple(
                    config=cfg,
                    checkpoint=checkpoint,
                    metadata=metadata,
                    parent_config=parent_cfg,
                )

        finally:
            cur.close()

    # ── Core interface: put ──────────────────────────────────────────

    def put(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        """Persist a checkpoint with its configuration and metadata.

        Uses HANA's ``UPSERT … WITH PRIMARY KEY`` for atomic
        insert-or-update semantics.
        """
        thread_id, ns = self._cfg(config)
        cp_id = checkpoint["id"]
        parent_id = config.get("configurable", {}).get("checkpoint_id")

        type_str, cp_blob = self._dump(checkpoint)
        meta_json = json.dumps(metadata) if metadata else None

        cur = self.conn.cursor()
        try:
            cur.execute(
                _UPSERT_CHECKPOINT,
                (
                    thread_id,
                    ns,
                    cp_id,
                    parent_id,
                    type_str,
                    cp_blob,
                    meta_json,
                ),
            )
            self.conn.commit()
        finally:
            cur.close()

        return {
            "configurable": {
                "thread_id": thread_id,
                "checkpoint_ns": ns,
                "checkpoint_id": cp_id,
            }
        }

    # ── Core interface: put_writes ───────────────────────────────────

    def put_writes(
        self,
        config: RunnableConfig,
        writes: Sequence[tuple[str, Any]],
        task_id: str,
        task_path: str = "",
    ) -> None:
        """Store intermediate writes linked to a checkpoint.

        These are the pending writes that allow LangGraph to skip
        already-completed nodes when resuming after a failure.
        """
        thread_id, ns = self._cfg(config)
        cp_id = config.get("configurable", {}).get("checkpoint_id", "")

        cur = self.conn.cursor()
        try:
            for idx, (channel, value) in enumerate(writes):
                if value is not None:
                    type_str, blob = self._dump(value)
                else:
                    type_str, blob = None, None

                cur.execute(
                    _UPSERT_WRITE,
                    (
                        thread_id,
                        ns,
                        cp_id,
                        task_id,
                        task_path,
                        idx,
                        channel,
                        type_str,
                        blob,
                    ),
                )
            self.conn.commit()
        finally:
            cur.close()

    # ── Core interface: delete_thread ────────────────────────────────

    def delete_thread(self, thread_id: str) -> None:
        """Delete all checkpoints and writes for a thread.

        This permanently removes a thread's entire history.  Useful for
        cleanup, GDPR compliance, or recovering from corrupted state
        (e.g. orphaned tool calls that cause downstream errors).
        """
        cur = self.conn.cursor()
        try:
            cur.execute(
                f'DELETE FROM "{_WRITES_TABLE}" WHERE thread_id = ?',
                (thread_id,),
            )
            cur.execute(
                f'DELETE FROM "{_CHECKPOINTS_TABLE}" WHERE thread_id = ?',
                (thread_id,),
            )
            self.conn.commit()
            logger.info("Deleted thread %s", thread_id)
        except Exception:
            logger.exception("Failed to delete thread %s", thread_id)
            raise
        finally:
            cur.close()

    # ── Async interface ──────────────────────────────────────────────
    # hdbcli is a synchronous C-extension driver.  We delegate async
    # methods to their sync counterparts.  For true async in high-
    # concurrency deployments, wrap calls with asyncio.to_thread().

    async def aget_tuple(
        self, config: RunnableConfig
    ) -> Optional[CheckpointTuple]:
        return self.get_tuple(config)

    async def alist(
        self,
        config: Optional[RunnableConfig],
        *,
        filter: Optional[dict[str, Any]] = None,
        before: Optional[RunnableConfig] = None,
        limit: Optional[int] = None,
    ) -> AsyncIterator[CheckpointTuple]:
        for item in self.list(
            config, filter=filter, before=before, limit=limit
        ):
            yield item

    async def aput(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        return self.put(config, checkpoint, metadata, new_versions)

    async def aput_writes(
        self,
        config: RunnableConfig,
        writes: Sequence[tuple[str, Any]],
        task_id: str,
        task_path: str = "",
    ) -> None:
        self.put_writes(config, writes, task_id, task_path)

    async def adelete_thread(self, thread_id: str) -> None:
        self.delete_thread(thread_id)
