import json
import logging
from contextlib import contextmanager, asynccontextmanager
from typing import (
    Any,
    AsyncIterator,
    Callable,
    Iterator,
    Optional,
    Sequence,
    Tuple,
    List,
    Dict,
)

import psycopg
from psycopg import Connection, AsyncConnection
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool, AsyncConnectionPool

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

logger = logging.getLogger(__name__)

class ContentDecoder:
    def __init__(self, serde: Any):
        self.serde = serde

    def decode(self, content: bytes, type_: str) -> Any:
        return self.serde.loads_typed((type_, content))

class TimescaleDBSaver(BaseCheckpointSaver[str]):
    pool: ConnectionPool

    def __init__(
        self,
        pool: ConnectionPool,
        *,
        serde: Optional[Any] = None,
    ) -> None:
        super().__init__(serde=serde or JsonPlusSerializer())
        self.pool = pool
        self._decoder = ContentDecoder(self.serde)

    @classmethod
    @contextmanager
    def from_conn_string(cls, conn_string: str, min_size: int = 4, max_size: int = 20, **kwargs: Any) -> Iterator["TimescaleDBSaver"]:
        with ConnectionPool(
            conn_string,
            min_size=min_size,
            max_size=max_size,
            kwargs={"autocommit": False, **kwargs},
        ) as pool:
            yield cls(pool=pool)

    def setup(self) -> None:
        """Create tables and hypertables in TimescaleDB."""
        with self.pool.connection() as conn:
            with conn.cursor() as cur:
                # Disable synchronous commit for maximum write throughput
                cur.execute("SET synchronous_commit = off;")

                cur.execute(
                    """
                    CREATE UNLOGGED TABLE IF NOT EXISTS checkpoints (
                        thread_id TEXT NOT NULL,
                        checkpoint_ns TEXT NOT NULL DEFAULT '',
                        checkpoint_id TEXT NOT NULL,
                        parent_checkpoint_id TEXT,
                        type TEXT,
                        checkpoint BYTEA NOT NULL,
                        metadata BYTEA NOT NULL,
                        timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW()
                    );
                    """
                )
                cur.execute(
                    """
                    CREATE INDEX IF NOT EXISTS checkpoints_thread_idx ON checkpoints (thread_id, checkpoint_ns, timestamp DESC);
                    """
                )

                cur.execute(
                    """
                    CREATE UNLOGGED TABLE IF NOT EXISTS checkpoint_writes (
                        thread_id TEXT NOT NULL,
                        checkpoint_ns TEXT NOT NULL DEFAULT '',
                        checkpoint_id TEXT NOT NULL,
                        task_id TEXT NOT NULL,
                        idx INTEGER NOT NULL,
                        channel TEXT NOT NULL,
                        type TEXT,
                        value BYTEA NOT NULL,
                        timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW()
                    );
                    """
                )
                cur.execute(
                    """
                    CREATE INDEX IF NOT EXISTS checkpoint_writes_thread_idx ON checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, timestamp ASC);
                    """
                )

            conn.commit()

    def _get_writes(self, conn: Connection, thread_id: str, checkpoint_ns: str, checkpoint_id: str) -> List[PendingWrite]:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute(
                """
                SELECT task_id, channel, type, value
                FROM checkpoint_writes
                WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
                ORDER BY timestamp ASC, idx ASC
                """,
                (thread_id, checkpoint_ns, checkpoint_id)
            )
            rows = cur.fetchall()
            writes = []
            for row in rows:
                writes.append((row["task_id"], row["channel"], self._decoder.decode(row["value"], row["type"] or "json")))
            return writes

    def get_tuple(self, config: dict) -> Optional[CheckpointTuple]:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"].get("checkpoint_id")

        with self.pool.connection() as conn:
            with conn.cursor(row_factory=dict_row) as cur:
                if checkpoint_id:
                    cur.execute(
                        """
                        SELECT checkpoint, metadata, parent_checkpoint_id, type
                        FROM checkpoints
                        WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
                        ORDER BY timestamp DESC
                        LIMIT 1
                        """,
                        (thread_id, checkpoint_ns, checkpoint_id)
                    )
                else:
                    cur.execute(
                        """
                        SELECT checkpoint, metadata, parent_checkpoint_id, checkpoint_id, type
                        FROM checkpoints
                        WHERE thread_id = %s AND checkpoint_ns = %s
                        ORDER BY timestamp DESC
                        LIMIT 1
                        """,
                        (thread_id, checkpoint_ns)
                    )
                row = cur.fetchone()
                
                if not row:
                    return None

                if not checkpoint_id:
                    checkpoint_id = row["checkpoint_id"]

                checkpoint = self._decoder.decode(row["checkpoint"], row["type"] or "json")
                metadata = json.loads(row["metadata"].decode("utf-8"))
                parent_config = None
                if row["parent_checkpoint_id"]:
                    parent_config = {
                        "configurable": {
                            "thread_id": thread_id,
                            "checkpoint_ns": checkpoint_ns,
                            "checkpoint_id": row["parent_checkpoint_id"],
                        }
                    }

                writes = self._get_writes(conn, thread_id, checkpoint_ns, checkpoint_id)

                return CheckpointTuple(
                    config={
                        "configurable": {
                            "thread_id": thread_id,
                            "checkpoint_ns": checkpoint_ns,
                            "checkpoint_id": checkpoint_id,
                        }
                    },
                    checkpoint=checkpoint,
                    metadata=metadata,
                    parent_config=parent_config,
                    pending_writes=writes,
                )

    def list(
        self,
        config: Optional[dict],
        *,
        filter: Optional[dict] = None,
        before: Optional[dict] = None,
        limit: Optional[int] = None,
    ) -> Iterator[CheckpointTuple]:
        thread_id = config["configurable"]["thread_id"] if config else None
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "") if config else None
        
        query = "SELECT checkpoint, metadata, parent_checkpoint_id, checkpoint_id, type FROM checkpoints WHERE 1=1"
        params = []
        
        if thread_id:
            query += " AND thread_id = %s"
            params.append(thread_id)
        if checkpoint_ns is not None:
            query += " AND checkpoint_ns = %s"
            params.append(checkpoint_ns)
            
        if before and "checkpoint_id" in before["configurable"]:
            query += " AND checkpoint_id < %s"
            params.append(before["configurable"]["checkpoint_id"])

        query += " ORDER BY timestamp DESC"
        if limit:
            query += " LIMIT %s"
            params.append(limit)

        with self.pool.connection() as conn:
            with conn.cursor(row_factory=dict_row) as cur:
                cur.execute(query, tuple(params))
                for row in cur.fetchall():
                    checkpoint_id = row["checkpoint_id"]
                    checkpoint = self._decoder.decode(row["checkpoint"], row["type"] or "json")
                    metadata = json.loads(row["metadata"].decode("utf-8"))
                    parent_config = None
                    if row["parent_checkpoint_id"]:
                        parent_config = {
                            "configurable": {
                                "thread_id": thread_id,
                                "checkpoint_ns": checkpoint_ns,
                                "checkpoint_id": row["parent_checkpoint_id"],
                            }
                        }

                    writes = self._get_writes(conn, thread_id, checkpoint_ns, checkpoint_id)

                    yield CheckpointTuple(
                        config={
                            "configurable": {
                                "thread_id": thread_id,
                                "checkpoint_ns": checkpoint_ns,
                                "checkpoint_id": checkpoint_id,
                            }
                        },
                        checkpoint=checkpoint,
                        metadata=metadata,
                        parent_config=parent_config,
                        pending_writes=writes,
                    )

    def put(
        self,
        config: dict,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> dict:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = checkpoint["id"]

        parent_checkpoint_id = config["configurable"].get("checkpoint_id")
        type_ = metadata.get("source", "input")

        # Serialize outside the connection context to minimise hold time
        checkpoint_type, checkpoint_blob = self.serde.dumps_typed(checkpoint)
        metadata_blob = json.dumps(metadata).encode("utf-8")

        with self.pool.connection() as conn:
            # Pipeline mode: SET + INSERT are sent in a single roundtrip
            with conn.pipeline():
                with conn.cursor() as cur:
                    cur.execute("SET synchronous_commit = off;")
                    cur.execute(
                        """
                        INSERT INTO checkpoints (
                            thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
                        ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                        """,
                        (
                            thread_id,
                            checkpoint_ns,
                            checkpoint_id,
                            parent_checkpoint_id,
                            checkpoint_type,
                            checkpoint_blob,
                            metadata_blob,
                        )
                    )
            conn.commit()

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

    def put_writes(
        self,
        config: dict,
        writes: Sequence[Tuple[str, Any]],
        task_id: str,
        task_path: str = "",
    ) -> None:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"]["checkpoint_id"]

        params = []
        for idx, (channel, value) in enumerate(writes):
            mapped_idx = WRITES_IDX_MAP.get(channel, idx)
            type_, value_blob = self.serde.dumps_typed(value)
            params.append((thread_id, checkpoint_ns, checkpoint_id, task_id, mapped_idx, channel, type_, value_blob))

        with self.pool.connection() as conn:
            with conn.cursor() as cur:
                cur.executemany(
                    """
                    INSERT INTO checkpoint_writes (
                        thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value, timestamp
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, now())
                    """,
                    params
                )
            conn.commit()


class AsyncTimescaleDBSaver(BaseCheckpointSaver[str]):
    pool: AsyncConnectionPool

    def __init__(
        self,
        pool: AsyncConnectionPool,
        *,
        serde: Optional[Any] = None,
    ) -> None:
        super().__init__(serde=serde or JsonPlusSerializer())
        self.pool = pool
        self._decoder = ContentDecoder(self.serde)

    @classmethod
    @asynccontextmanager
    async def from_conn_string(cls, conn_string: str, **kwargs: Any) -> AsyncIterator["AsyncTimescaleDBSaver"]:
        async with AsyncConnectionPool(conn_string, kwargs=kwargs) as pool:
            yield cls(pool=pool)

    async def setup(self) -> None:
        """Create tables and hypertables in TimescaleDB."""
        async with self.pool.connection() as conn:
            async with conn.cursor() as cur:
                await cur.execute(
                    """
                    CREATE TABLE IF NOT EXISTS checkpoints (
                        thread_id TEXT NOT NULL,
                        checkpoint_ns TEXT NOT NULL DEFAULT '',
                        checkpoint_id TEXT NOT NULL,
                        parent_checkpoint_id TEXT,
                        type TEXT,
                        checkpoint BYTEA NOT NULL,
                        metadata BYTEA NOT NULL,
                        timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW()
                    );
                    """
                )
                await cur.execute(
                    """
                    CREATE UNIQUE INDEX IF NOT EXISTS checkpoints_pkey ON checkpoints (thread_id, checkpoint_ns, checkpoint_id, timestamp);
                    """
                )
                await cur.execute(
                    """
                    SELECT count(*) FROM _timescaledb_catalog.hypertable 
                    WHERE table_name = 'checkpoints';
                    """
                )
                row = await cur.fetchone()
                if row and row[0] == 0:
                    await cur.execute(
                        "SELECT create_hypertable('checkpoints', 'timestamp', migrate_data => true);"
                    )

                await cur.execute(
                    """
                    CREATE TABLE IF NOT EXISTS checkpoint_writes (
                        thread_id TEXT NOT NULL,
                        checkpoint_ns TEXT NOT NULL DEFAULT '',
                        checkpoint_id TEXT NOT NULL,
                        task_id TEXT NOT NULL,
                        idx INTEGER NOT NULL,
                        channel TEXT NOT NULL,
                        type TEXT,
                        value BYTEA NOT NULL,
                        timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW()
                    );
                    """
                )
                await cur.execute(
                    """
                    CREATE UNIQUE INDEX IF NOT EXISTS checkpoint_writes_pkey ON checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, timestamp);
                    """
                )
                await cur.execute(
                    """
                    SELECT count(*) FROM _timescaledb_catalog.hypertable 
                    WHERE table_name = 'checkpoint_writes';
                    """
                )
                row = await cur.fetchone()
                if row and row[0] == 0:
                    await cur.execute(
                        "SELECT create_hypertable('checkpoint_writes', 'timestamp', migrate_data => true);"
                    )

            await conn.commit()

    async def _aget_writes(self, conn: AsyncConnection, thread_id: str, checkpoint_ns: str, checkpoint_id: str) -> List[PendingWrite]:
        async with conn.cursor(row_factory=dict_row) as cur:
            await cur.execute(
                """
                SELECT task_id, channel, type, value
                FROM checkpoint_writes
                WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
                ORDER BY timestamp ASC, idx ASC
                """,
                (thread_id, checkpoint_ns, checkpoint_id)
            )
            rows = await cur.fetchall()
            writes = []
            for row in rows:
                writes.append((row["task_id"], row["channel"], self._decoder.decode(row["value"], row["type"] or "json")))
            return writes

    async def aget_tuple(self, config: dict) -> Optional[CheckpointTuple]:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"].get("checkpoint_id")

        async with self.pool.connection() as conn:
            async with conn.cursor(row_factory=dict_row) as cur:
                if checkpoint_id:
                    await cur.execute(
                        """
                        SELECT checkpoint, metadata, parent_checkpoint_id, type
                        FROM checkpoints
                        WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
                        ORDER BY timestamp DESC
                        LIMIT 1
                        """,
                        (thread_id, checkpoint_ns, checkpoint_id)
                    )
                else:
                    await cur.execute(
                        """
                        SELECT checkpoint, metadata, parent_checkpoint_id, checkpoint_id, type
                        FROM checkpoints
                        WHERE thread_id = %s AND checkpoint_ns = %s
                        ORDER BY timestamp DESC
                        LIMIT 1
                        """,
                        (thread_id, checkpoint_ns)
                    )
                row = await cur.fetchone()
                
                if not row:
                    return None

                if not checkpoint_id:
                    checkpoint_id = row["checkpoint_id"]

                checkpoint = self._decoder.decode(row["checkpoint"], row["type"] or "json")
                metadata = json.loads(row["metadata"].decode("utf-8"))
                parent_config = None
                if row["parent_checkpoint_id"]:
                    parent_config = {
                        "configurable": {
                            "thread_id": thread_id,
                            "checkpoint_ns": checkpoint_ns,
                            "checkpoint_id": row["parent_checkpoint_id"],
                        }
                    }

                writes = await self._aget_writes(conn, thread_id, checkpoint_ns, checkpoint_id)

                return CheckpointTuple(
                    config={
                        "configurable": {
                            "thread_id": thread_id,
                            "checkpoint_ns": checkpoint_ns,
                            "checkpoint_id": checkpoint_id,
                        }
                    },
                    checkpoint=checkpoint,
                    metadata=metadata,
                    parent_config=parent_config,
                    pending_writes=writes,
                )

    async def alist(
        self,
        config: Optional[dict],
        *,
        filter: Optional[dict] = None,
        before: Optional[dict] = None,
        limit: Optional[int] = None,
    ) -> AsyncIterator[CheckpointTuple]:
        thread_id = config["configurable"]["thread_id"] if config else None
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "") if config else None
        
        query = "SELECT checkpoint, metadata, parent_checkpoint_id, checkpoint_id FROM checkpoints WHERE 1=1"
        params = []
        
        if thread_id:
            query += " AND thread_id = %s"
            params.append(thread_id)
        if checkpoint_ns is not None:
            query += " AND checkpoint_ns = %s"
            params.append(checkpoint_ns)
            
        if before and "checkpoint_id" in before["configurable"]:
            query += " AND checkpoint_id < %s"
            params.append(before["configurable"]["checkpoint_id"])

        query += " ORDER BY timestamp DESC"
        if limit:
            query += " LIMIT %s"
            params.append(limit)

        async with self.pool.connection() as conn:
            async with conn.cursor(row_factory=dict_row) as cur:
                await cur.execute(query, tuple(params))
                for row in await cur.fetchall():
                    checkpoint_id = row["checkpoint_id"]
                    checkpoint = self._decoder.decode(row["checkpoint"], row["type"] or "json")
                    metadata = json.loads(row["metadata"].decode("utf-8"))
                    parent_config = None
                    if row["parent_checkpoint_id"]:
                        parent_config = {
                            "configurable": {
                                "thread_id": thread_id,
                                "checkpoint_ns": checkpoint_ns,
                                "checkpoint_id": row["parent_checkpoint_id"],
                            }
                        }

                    writes = await self._aget_writes(conn, thread_id, checkpoint_ns, checkpoint_id)

                    yield CheckpointTuple(
                        config={
                            "configurable": {
                                "thread_id": thread_id,
                                "checkpoint_ns": checkpoint_ns,
                                "checkpoint_id": checkpoint_id,
                            }
                        },
                        checkpoint=checkpoint,
                        metadata=metadata,
                        parent_config=parent_config,
                        pending_writes=writes,
                    )

    async def aput(
        self,
        config: dict,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> dict:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = checkpoint["id"]
        
        parent_checkpoint_id = config["configurable"].get("checkpoint_id")
        type_ = metadata.get("source", "input")
        
        checkpoint_type, checkpoint_blob = self.serde.dumps_typed(checkpoint)
        metadata_blob = json.dumps(metadata).encode("utf-8")

        async with self.pool.connection() as conn:
            async with conn.cursor() as cur:
                await cur.execute(
                    """
                    INSERT INTO checkpoints (
                        thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        thread_id,
                        checkpoint_ns,
                        checkpoint_id,
                        parent_checkpoint_id,
                        checkpoint_type,
                        checkpoint_blob,
                        metadata_blob,
                    )
                )
            await conn.commit()

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

    async def aput_writes(
        self,
        config: dict,
        writes: Sequence[Tuple[str, Any]],
        task_id: str,
        task_path: str = "",
    ) -> None:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
        checkpoint_id = config["configurable"]["checkpoint_id"]

        params = []
        for idx, (channel, value) in enumerate(writes):
            mapped_idx = WRITES_IDX_MAP.get(channel, idx)
            type_, value_blob = self.serde.dumps_typed(value)
            params.append((thread_id, checkpoint_ns, checkpoint_id, task_id, mapped_idx, channel, type_, value_blob))

        async with self.pool.connection() as conn:
            async with conn.cursor() as cur:
                await cur.executemany(
                    """
                    INSERT INTO checkpoint_writes (
                        thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value, timestamp
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, now())
                    """,
                    params
                )
            await conn.commit()
