"""LangGraph checkpoint saver backed by Oracle ADB."""

from __future__ import annotations

from collections.abc import Iterator, Mapping, Sequence
from copy import deepcopy
import logging
from random import random
from typing import Any

from langgraph.checkpoint.base import (
    BaseCheckpointSaver,
    ChannelVersions,
    Checkpoint,
    CheckpointMetadata,
    CheckpointTuple,
    WRITES_IDX_MAP,
    get_checkpoint_id,
    get_serializable_checkpoint_metadata,
)
from langgraph.checkpoint.serde.types import _DeltaSnapshot

from agente_contas_tim.repositories.oracle_state_connection import (
    OracleStateConnectionFactory,
    is_unique_constraint_error,
    normalize_json,
    set_blob_inputsizes,
    read_json_value,
    set_json_inputsizes,
)


_DEFAULT_NS = "__default__"
_PRIMITIVE_TYPES = (str, int, float, bool)
logger = logging.getLogger(__name__)


class OracleCheckpointSaver(BaseCheckpointSaver):
    """Synchronous LangGraph checkpointer persisted in Oracle native JSON/BLOB."""

    _REQUIRED_TABLES = {
        "TB_CONT_WORKFLOW_CHECKPOINT",
        "TB_CONT_WORKFLOW_CHECKPOINT_WRITE",
        "TB_CONT_WORKFLOW_CHECKPOINT_BLOB",
    }

    def __init__(self, connection_factory: OracleStateConnectionFactory) -> None:
        super().__init__()
        self._connection_factory = connection_factory
        self._connection_factory.validate_required_tables(self._REQUIRED_TABLES)

    def setup(self) -> None:
        self._connection_factory.validate_required_tables(self._REQUIRED_TABLES)

    def get_tuple(self, config: Mapping[str, Any]) -> CheckpointTuple | None:
        configurable = dict(config.get("configurable", {}) or {})
        thread_id = str(configurable["thread_id"])
        checkpoint_ns = _to_db_ns(str(configurable.get("checkpoint_ns", "")))
        checkpoint_id = get_checkpoint_id(config)  # type: ignore[arg-type]

        with self._connection_factory.connect() as conn, conn.cursor() as cur:
            if checkpoint_id:
                cur.execute(
                    """
                    SELECT
                      thread_id,
                      checkpoint_ns,
                      checkpoint_id,
                      parent_checkpoint_id,
                      checkpoint_json,
                      metadata_json
                    FROM tb_cont_workflow_checkpoint
                    WHERE thread_id = :thread_id
                      AND (
                        checkpoint_ns = :checkpoint_ns
                        OR (
                          :checkpoint_ns = :default_checkpoint_ns
                          AND checkpoint_ns IS NULL
                        )
                      )
                      AND checkpoint_id = :checkpoint_id
                    """,
                    {
                        "thread_id": thread_id,
                        "checkpoint_ns": checkpoint_ns,
                        "default_checkpoint_ns": _DEFAULT_NS,
                        "checkpoint_id": checkpoint_id,
                    },
                )
            else:
                cur.execute(
                    """
                    SELECT
                      thread_id,
                      checkpoint_ns,
                      checkpoint_id,
                      parent_checkpoint_id,
                      checkpoint_json,
                      metadata_json
                    FROM tb_cont_workflow_checkpoint
                    WHERE thread_id = :thread_id
                      AND (
                        checkpoint_ns = :checkpoint_ns
                        OR (
                          :checkpoint_ns = :default_checkpoint_ns
                          AND checkpoint_ns IS NULL
                        )
                      )
                    ORDER BY checkpoint_id DESC
                    FETCH FIRST 1 ROWS ONLY
                    """,
                    {
                        "thread_id": thread_id,
                        "checkpoint_ns": checkpoint_ns,
                        "default_checkpoint_ns": _DEFAULT_NS,
                    },
                )
            row = _fetchone_dict(cur)
            if row is None:
                return None
            return self._load_checkpoint_tuple(cur, row)

    def list(
        self,
        config: Mapping[str, Any] | None,
        *,
        filter: dict[str, Any] | None = None,
        before: Mapping[str, Any] | None = None,
        limit: int | None = None,
    ) -> Iterator[CheckpointTuple]:
        clauses: list[str] = []
        params: dict[str, Any] = {}
        configurable = dict((config or {}).get("configurable", {}) or {})
        if configurable.get("thread_id"):
            clauses.append("thread_id = :thread_id")
            params["thread_id"] = str(configurable["thread_id"])
        if "checkpoint_ns" in configurable:
            clauses.append(
                "(checkpoint_ns = :checkpoint_ns OR "
                "(:checkpoint_ns = :default_checkpoint_ns "
                "AND checkpoint_ns IS NULL))"
            )
            params["checkpoint_ns"] = _to_db_ns(
                str(configurable.get("checkpoint_ns", ""))
            )
            params["default_checkpoint_ns"] = _DEFAULT_NS
        if before:
            before_id = get_checkpoint_id(before)  # type: ignore[arg-type]
            if before_id:
                clauses.append("checkpoint_id < :before_checkpoint_id")
                params["before_checkpoint_id"] = before_id

        where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
        query = f"""
            SELECT
              thread_id,
              checkpoint_ns,
              checkpoint_id,
              parent_checkpoint_id,
              checkpoint_json,
              metadata_json
            FROM tb_cont_workflow_checkpoint
            {where}
            ORDER BY checkpoint_id DESC
        """

        emitted = 0
        with self._connection_factory.connect() as conn, conn.cursor() as cur:
            cur.execute(query, params)
            rows = [_row_to_dict(cur, row) for row in cur.fetchall()]
            for row in rows:
                metadata = read_json_value(row.get("metadata_json")) or {}
                if filter and not _metadata_matches(metadata, filter):
                    continue
                yield self._load_checkpoint_tuple(cur, row)
                emitted += 1
                if limit is not None and emitted >= limit:
                    return

    def put(
        self,
        config: Mapping[str, Any],
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> dict[str, Any]:
        configurable = dict(config["configurable"])
        thread_id = str(configurable.pop("thread_id"))
        checkpoint_ns = str(configurable.pop("checkpoint_ns", ""))
        parent_checkpoint_id = configurable.pop("checkpoint_id", None)
        checkpoint_id = str(checkpoint["id"])
        db_checkpoint_ns = _to_db_ns(checkpoint_ns)

        checkpoint_copy = deepcopy(checkpoint)
        channel_values = dict(checkpoint_copy.get("channel_values") or {})
        checkpoint_copy["channel_values"] = channel_values
        blob_values: dict[str, Any] = {}
        for channel, value in list(channel_values.items()):
            if isinstance(value, _DeltaSnapshot):
                blob_values[channel] = value
                channel_values[channel] = True
                continue
            if value is None or isinstance(value, _PRIMITIVE_TYPES):
                continue
            blob_values[channel] = channel_values.pop(channel)

        try:
            with self._connection_factory.connect() as conn, conn.cursor() as cur:
                blob_versions = {
                    channel: version
                    for channel, version in dict(new_versions).items()
                    if channel in blob_values
                }
                for channel, version in blob_versions.items():
                    self._upsert_blob(
                        cur,
                        thread_id=thread_id,
                        checkpoint_ns=db_checkpoint_ns,
                        channel=str(channel),
                        version=str(version),
                        value=blob_values[channel],
                    )

                set_json_inputsizes(cur, "checkpoint_json", "metadata_json")
                cur.execute(
                    """
                    MERGE INTO tb_cont_workflow_checkpoint dst
                    USING (
                      SELECT
                        :thread_id thread_id,
                        :checkpoint_ns checkpoint_ns,
                        :checkpoint_id checkpoint_id
                      FROM dual
                    ) src
                       ON (
                         dst.thread_id = src.thread_id
                         AND dst.checkpoint_ns = src.checkpoint_ns
                         AND dst.checkpoint_id = src.checkpoint_id
                       )
                    WHEN MATCHED THEN UPDATE SET
                      parent_checkpoint_id = :parent_checkpoint_id,
                      checkpoint_json = :checkpoint_json,
                      metadata_json = :metadata_json
                    WHEN NOT MATCHED THEN INSERT (
                      thread_id,
                      checkpoint_ns,
                      checkpoint_id,
                      parent_checkpoint_id,
                      checkpoint_json,
                      metadata_json
                    ) VALUES (
                      :thread_id,
                      :checkpoint_ns,
                      :checkpoint_id,
                      :parent_checkpoint_id,
                      :checkpoint_json,
                      :metadata_json
                    )
                    """,
                    {
                        "thread_id": thread_id,
                        "checkpoint_ns": db_checkpoint_ns,
                        "checkpoint_id": checkpoint_id,
                        "parent_checkpoint_id": parent_checkpoint_id,
                        "checkpoint_json": normalize_json(checkpoint_copy),
                        "metadata_json": normalize_json(
                            get_serializable_checkpoint_metadata(
                                config,  # type: ignore[arg-type]
                                metadata,
                            )
                        ),
                    },
                )
        except Exception:
            logger.error(
                "oracle_checkpoint.put.failed thread_id=%s checkpoint_ns=%s "
                "checkpoint_id=%s parent_checkpoint_id=%s blob_channels=%s "
                "channel_keys=%s",
                thread_id,
                db_checkpoint_ns,
                checkpoint_id,
                parent_checkpoint_id,
                sorted(str(channel) for channel in blob_values.keys()),
                sorted(str(channel) for channel in channel_values.keys()),
                exc_info=True,
            )
            raise

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

    def put_writes(
        self,
        config: Mapping[str, Any],
        writes: Sequence[tuple[str, Any]],
        task_id: str,
        task_path: str = "",
    ) -> None:
        configurable = dict(config["configurable"])
        thread_id = str(configurable["thread_id"])
        checkpoint_ns = _to_db_ns(str(configurable.get("checkpoint_ns", "")))
        checkpoint_id = str(configurable["checkpoint_id"])
        use_upsert = all(channel in WRITES_IDX_MAP for channel, _ in writes)

        try:
            with self._connection_factory.connect() as conn, conn.cursor() as cur:
                set_blob_inputsizes(cur, "blob_payload")
                for idx, (channel, value) in enumerate(writes):
                    write_idx = WRITES_IDX_MAP.get(channel, idx)
                    type_name, payload = self.serde.dumps_typed(value)
                    params = {
                        "thread_id": thread_id,
                        "checkpoint_ns": checkpoint_ns,
                        "checkpoint_id": checkpoint_id,
                        "task_id": task_id,
                        "task_path": task_path,
                        "idx": write_idx,
                        "channel": channel,
                        "type_name": type_name,
                        "blob_payload": _blob_payload_for_storage(
                            type_name,
                            payload,
                        ),
                    }
                    if use_upsert:
                        cur.execute(
                            """
                            MERGE INTO tb_cont_workflow_checkpoint_write dst
                            USING (
                              SELECT
                                :thread_id thread_id,
                                :checkpoint_ns checkpoint_ns,
                                :checkpoint_id checkpoint_id,
                                :task_id task_id,
                                :idx idx
                              FROM dual
                            ) src
                               ON (
                                 dst.thread_id = src.thread_id
                                 AND dst.checkpoint_ns = src.checkpoint_ns
                                 AND dst.checkpoint_id = src.checkpoint_id
                                 AND dst.task_id = src.task_id
                                 AND dst.idx = src.idx
                               )
                            WHEN MATCHED THEN UPDATE SET
                              task_path = :task_path,
                              channel = :channel,
                              type_name = :type_name,
                              blob_payload = :blob_payload,
                              value_json = NULL
                            WHEN NOT MATCHED THEN INSERT (
                              thread_id,
                              checkpoint_ns,
                              checkpoint_id,
                              task_id,
                              task_path,
                              idx,
                              channel,
                              type_name,
                              blob_payload,
                              value_json
                            ) VALUES (
                              :thread_id,
                              :checkpoint_ns,
                              :checkpoint_id,
                              :task_id,
                              :task_path,
                              :idx,
                              :channel,
                              :type_name,
                              :blob_payload,
                              NULL
                            )
                            """,
                            params,
                        )
                    else:
                        try:
                            cur.execute(
                                """
                                INSERT INTO tb_cont_workflow_checkpoint_write (
                                  thread_id,
                                  checkpoint_ns,
                                  checkpoint_id,
                                  task_id,
                                  task_path,
                                  idx,
                                  channel,
                                  type_name,
                                  blob_payload,
                                  value_json
                                ) VALUES (
                                  :thread_id,
                                  :checkpoint_ns,
                                  :checkpoint_id,
                                  :task_id,
                                  :task_path,
                                  :idx,
                                  :channel,
                                  :type_name,
                                  :blob_payload,
                                  NULL
                                )
                                """,
                                params,
                            )
                        except Exception as exc:
                            if is_unique_constraint_error(exc):
                                continue
                            raise
        except Exception:
            logger.error(
                "oracle_checkpoint.put_writes.failed thread_id=%s "
                "checkpoint_ns=%s checkpoint_id=%s task_id=%s "
                "task_path=%s channels=%s use_upsert=%s",
                thread_id,
                checkpoint_ns,
                checkpoint_id,
                task_id,
                task_path,
                [str(channel) for channel, _ in writes],
                use_upsert,
                exc_info=True,
            )
            raise

    def delete_thread(self, thread_id: str) -> None:
        with self._connection_factory.connect() as conn, conn.cursor() as cur:
            cur.execute(
                "DELETE FROM tb_cont_workflow_checkpoint_write WHERE thread_id = :thread_id",
                {"thread_id": thread_id},
            )
            cur.execute(
                "DELETE FROM tb_cont_workflow_checkpoint_blob WHERE thread_id = :thread_id",
                {"thread_id": thread_id},
            )
            cur.execute(
                "DELETE FROM tb_cont_workflow_checkpoint WHERE thread_id = :thread_id",
                {"thread_id": thread_id},
            )

    def close(self) -> None:
        return None

    def _load_checkpoint_tuple(
        self,
        cur: Any,
        row: dict[str, Any],
    ) -> CheckpointTuple:
        checkpoint = read_json_value(row.get("checkpoint_json")) or {}
        metadata = read_json_value(row.get("metadata_json")) or {}
        thread_id = str(row["thread_id"])
        checkpoint_ns = _db_ns_from_row(row.get("checkpoint_ns"))
        checkpoint_id = str(row["checkpoint_id"])
        parent_checkpoint_id = row.get("parent_checkpoint_id")
        channel_values = dict(checkpoint.get("channel_values") or {})
        checkpoint["channel_values"] = {
            **channel_values,
            **self._load_blobs(
                cur,
                checkpoint=checkpoint,
                thread_id=thread_id,
                checkpoint_ns=checkpoint_ns,
            ),
        }
        pending_writes = self._load_writes(
            cur,
            thread_id=thread_id,
            checkpoint_ns=checkpoint_ns,
            checkpoint_id=checkpoint_id,
        )
        return CheckpointTuple(
            {
                "configurable": {
                    "thread_id": thread_id,
                    "checkpoint_ns": _from_db_ns(checkpoint_ns),
                    "checkpoint_id": checkpoint_id,
                }
            },
            checkpoint,
            metadata,
            (
                {
                    "configurable": {
                        "thread_id": thread_id,
                        "checkpoint_ns": _from_db_ns(checkpoint_ns),
                        "checkpoint_id": parent_checkpoint_id,
                    }
                }
                if parent_checkpoint_id
                else None
            ),
            pending_writes,
        )

    def _load_blobs(
        self,
        cur: Any,
        *,
        checkpoint: dict[str, Any],
        thread_id: str,
        checkpoint_ns: str,
    ) -> dict[str, Any]:
        channel_versions = checkpoint.get("channel_versions") or {}
        if not isinstance(channel_versions, dict):
            return {}
        values: dict[str, Any] = {}
        for channel, version in channel_versions.items():
            cur.execute(
                """
                SELECT channel, type_name, blob_payload
                FROM tb_cont_workflow_checkpoint_blob
                WHERE thread_id = :thread_id
                  AND (
                    checkpoint_ns = :checkpoint_ns
                    OR (
                      :checkpoint_ns = :default_checkpoint_ns
                      AND checkpoint_ns IS NULL
                    )
                  )
                  AND channel = :channel
                  AND version = :version
                """,
                {
                    "thread_id": thread_id,
                    "checkpoint_ns": checkpoint_ns,
                    "default_checkpoint_ns": _DEFAULT_NS,
                    "channel": str(channel),
                    "version": str(version),
                },
            )
            blob_row = _fetchone_dict(cur)
            if blob_row is None:
                continue
            type_name = str(blob_row.get("type_name") or "")
            if type_name == "empty":
                continue
            payload = _read_lob_bytes(blob_row.get("blob_payload"))
            try:
                values[str(blob_row["channel"])] = self.serde.loads_typed(
                    (type_name, payload)
                )
            except Exception:
                logger.error(
                    "oracle_checkpoint.load_blob.failed thread_id=%s "
                    "checkpoint_ns=%s channel=%s version=%s type_name=%s "
                    "payload_len=%s",
                    thread_id,
                    checkpoint_ns,
                    channel,
                    version,
                    type_name,
                    len(payload),
                    exc_info=True,
                )
                raise
        return values

    def _load_writes(
        self,
        cur: Any,
        *,
        thread_id: str,
        checkpoint_ns: str,
        checkpoint_id: str,
    ) -> list[tuple[str, str, Any]]:
        cur.execute(
            """
            SELECT task_id, channel, type_name, blob_payload
            FROM tb_cont_workflow_checkpoint_write
            WHERE thread_id = :thread_id
              AND (
                checkpoint_ns = :checkpoint_ns
                OR (
                  :checkpoint_ns = :default_checkpoint_ns
                  AND checkpoint_ns IS NULL
                )
              )
              AND checkpoint_id = :checkpoint_id
            ORDER BY task_id, idx
            """,
            {
                "thread_id": thread_id,
                "checkpoint_ns": checkpoint_ns,
                "default_checkpoint_ns": _DEFAULT_NS,
                "checkpoint_id": checkpoint_id,
            },
        )
        rows = [_row_to_dict(cur, row) for row in cur.fetchall()]
        writes: list[tuple[str, str, Any]] = []
        for row in rows:
            type_name = str(row["type_name"])
            payload = _read_lob_bytes(row.get("blob_payload"))
            try:
                value = self.serde.loads_typed((type_name, payload))
            except Exception:
                logger.error(
                    "oracle_checkpoint.load_write.failed thread_id=%s "
                    "checkpoint_ns=%s checkpoint_id=%s task_id=%s "
                    "channel=%s type_name=%s payload_len=%s",
                    thread_id,
                    checkpoint_ns,
                    checkpoint_id,
                    row.get("task_id"),
                    row.get("channel"),
                    type_name,
                    len(payload),
                    exc_info=True,
                )
                raise
            writes.append(
                (
                    str(row["task_id"]),
                    str(row["channel"]),
                    value,
                )
            )
        return writes

    def _upsert_blob(
        self,
        cur: Any,
        *,
        thread_id: str,
        checkpoint_ns: str,
        channel: str,
        version: str,
        value: Any,
    ) -> None:
        type_name, payload = self.serde.dumps_typed(value)
        set_blob_inputsizes(cur, "blob_payload")
        cur.execute(
            """
            MERGE INTO tb_cont_workflow_checkpoint_blob dst
            USING (
              SELECT
                :thread_id thread_id,
                :checkpoint_ns checkpoint_ns,
                :channel channel,
                :version version
              FROM dual
            ) src
               ON (
                 dst.thread_id = src.thread_id
                 AND dst.checkpoint_ns = src.checkpoint_ns
                 AND dst.channel = src.channel
                 AND dst.version = src.version
               )
            WHEN MATCHED THEN UPDATE SET
              type_name = :type_name,
              blob_payload = :blob_payload,
              json_payload = NULL
            WHEN NOT MATCHED THEN INSERT (
              thread_id,
              checkpoint_ns,
              channel,
              version,
              type_name,
              blob_payload,
              json_payload
            ) VALUES (
              :thread_id,
              :checkpoint_ns,
              :channel,
              :version,
              :type_name,
              :blob_payload,
              NULL
            )
            """,
            {
                "thread_id": thread_id,
                "checkpoint_ns": checkpoint_ns,
                "channel": channel,
                "version": version,
                "type_name": type_name,
                "blob_payload": _blob_payload_for_storage(type_name, payload),
            },
        )

    def get_next_version(self, current: Any, channel: None) -> str:
        if current is None:
            current_v = 0
        elif isinstance(current, int):
            current_v = current
        else:
            current_v = int(str(current).split(".")[0])
        next_v = current_v + 1
        next_h = random()
        return f"{next_v:032}.{next_h:016}"


def _to_db_ns(checkpoint_ns: str) -> str:
    return checkpoint_ns or _DEFAULT_NS


def _from_db_ns(checkpoint_ns: str | None) -> str:
    return (
        ""
        if not checkpoint_ns or checkpoint_ns == _DEFAULT_NS
        else checkpoint_ns
    )


def _db_ns_from_row(checkpoint_ns: Any) -> str:
    if checkpoint_ns is None:
        return _DEFAULT_NS
    return str(checkpoint_ns) or _DEFAULT_NS


def _fetchone_dict(cur: Any) -> dict[str, Any] | None:
    row = cur.fetchone()
    if row is None:
        return None
    return _row_to_dict(cur, row)


def _row_to_dict(cur: Any, row: Any) -> dict[str, Any]:
    names = [str(col[0]).lower() for col in cur.description]
    return dict(zip(names, row, strict=False))


def _metadata_matches(metadata: Any, expected: dict[str, Any]) -> bool:
    if not isinstance(metadata, dict):
        return False
    for key, value in expected.items():
        if metadata.get(key) != value:
            return False
    return True


def _read_lob_bytes(value: Any) -> bytes:
    if value is None:
        return b""
    if isinstance(value, bytes):
        return value
    if isinstance(value, bytearray):
        return bytes(value)
    if isinstance(value, memoryview):
        return value.tobytes()
    if isinstance(value, str):
        return value.encode("utf-8")
    read = getattr(value, "read", None)
    if callable(read):
        return _read_lob_bytes(read())
    return bytes(value)


def _blob_payload_for_storage(type_name: str, payload: bytes | None) -> bytes:
    if payload:
        return payload
    if type_name in {"empty", "null"}:
        # Oracle can normalize zero-length BLOB binds to NULL. LangGraph's
        # serde ignores the payload for these tags, so a one-byte sentinel keeps
        # legacy payload checks and Oracle's empty-BLOB behavior from breaking
        # resume writes that legitimately carry None.
        return b"\x00"
    return payload if payload is not None else b""