"""
PostgreSQL target for CocoIndex.

This module provides a two-level target state system for PostgreSQL:
1. Table level: Creates/drops tables in the database
2. Row level: Upserts/deletes rows within tables
"""

from __future__ import annotations

import asyncio
import datetime
import decimal
import ipaddress
import inspect
import json
import uuid
from dataclasses import dataclass
from typing import (
    Any,
    Callable,
    Collection,
    Generic,
    Literal,
    NamedTuple,
    Sequence,
)

from typing_extensions import TypeVar

try:
    import asyncpg  # type: ignore
    import pgvector.asyncpg  # type: ignore
except ImportError as e:
    raise ImportError(
        "asyncpg and pgvector are required to use the PostgreSQL connector. Please install cocoindex[postgres]."
    ) from e

import numpy as np

import cocoindex as coco
from cocoindex.connectorkits import statediff
from cocoindex.connectorkits.fingerprint import fingerprint_object
from cocoindex._internal.datatype import (
    AnyType,
    MappingType,
    SequenceType,
    RecordType,
    TypeChecker,
    UnionType,
    analyze_type_info,
    is_record_type,
)
from cocoindex.resources import schema as res_schema
from cocoindex._internal.serde import unpickle_safe
from cocoindex._internal.context_keys import ContextKey, ContextProvider

# Type aliases
_RowKey = tuple[Any, ...]  # Primary key values as tuple
_ROW_KEY_CHECKER = TypeChecker(tuple[Any, ...])
_RowValue = dict[str, Any]  # Column name -> value
_RowFingerprint = bytes
ValueEncoder = Callable[[Any], Any]

# asyncpg enforces a protocol limit of 32767 bind parameters per query.
_BIND_LIMIT: int = 32767


def _qualified_table_name(table_name: str, pg_schema_name: str | None) -> str:
    """Return a properly quoted (optionally schema-qualified) table name."""

    if pg_schema_name:
        return f'"{pg_schema_name}"."{table_name}"'
    return f'"{table_name}"'


class PgType(NamedTuple):
    """
    Annotation to specify a PostgreSQL column type.

    Use with `typing.Annotated` to override the default type mapping:

    ```python
    from typing import Annotated
    from dataclasses import dataclass
    from cocoindex.connectors.postgres import PgType

    @dataclass
    class MyRow:
        # Use integer instead of default bigint
        id: Annotated[int, PgType("integer")]
        # Use real instead of default double precision
        value: Annotated[float, PgType("real")]
        # Use timestamp without timezone
        created_at: Annotated[datetime.datetime, PgType("timestamp")]
    ```
    """

    pg_type: str
    encoder: ValueEncoder | None = None


def _json_encoder(value: Any) -> str:
    """Encode a value to JSON string for asyncpg."""
    return json.dumps(value, default=str)


_PGVECTOR_TYPE_BASES: frozenset[str] = frozenset({"vector", "halfvec"})
_PGVECTOR_TYPE_PREFIXES: frozenset[str] = frozenset(
    {f"{base}(" for base in _PGVECTOR_TYPE_BASES}
)


def _is_pgvector_pg_type(pg_type: str) -> bool:
    """
    Return True if `pg_type` is a pgvector type (`vector(n)`, ...).

    This is used for extension checks and validation.
    """
    t = pg_type.lower().strip()
    return any(t.startswith(p) for p in _PGVECTOR_TYPE_PREFIXES)


class _TypeMapping(NamedTuple):
    """Mapping from Python type to PostgreSQL type with optional encoder."""

    pg_type: str
    encoder: ValueEncoder | None = None


# Global mapping for leaf types
# Based on asyncpg's type conversion: https://magicstack.github.io/asyncpg/current/usage.html#type-conversion
# For types that map to multiple PostgreSQL types, uses the broader one.
_LEAF_TYPE_MAPPINGS: dict[type, _TypeMapping] = {
    # Boolean
    bool: _TypeMapping("boolean"),
    # Numeric types (use broader types)
    int: _TypeMapping("bigint"),
    float: _TypeMapping("double precision"),
    decimal.Decimal: _TypeMapping("numeric"),
    # NumPy scalar integer types (finer-grained)
    np.int8: _TypeMapping("smallint"),
    np.int16: _TypeMapping("smallint"),
    np.int32: _TypeMapping("integer"),
    np.int64: _TypeMapping("bigint"),
    # NumPy scalar unsigned integer types (Postgres has no unsigned ints)
    np.uint8: _TypeMapping("smallint"),  # always fits
    np.uint16: _TypeMapping("integer"),  # can exceed smallint
    np.uint32: _TypeMapping("bigint"),  # can exceed integer
    np.uint64: _TypeMapping("numeric"),  # can exceed bigint
    # Platform-dependent aliases
    np.int_: _TypeMapping("bigint"),
    np.uint: _TypeMapping("numeric"),
    # NumPy scalar float types (finer-grained)
    np.float16: _TypeMapping("real"),
    np.float32: _TypeMapping("real"),
    np.float64: _TypeMapping("double precision"),
    # String types
    str: _TypeMapping("text"),
    bytes: _TypeMapping("bytea"),
    # UUID
    uuid.UUID: _TypeMapping("uuid"),
    # Date/time types (use timezone-aware variants as broader)
    datetime.date: _TypeMapping("date"),
    datetime.time: _TypeMapping("time with time zone"),
    datetime.datetime: _TypeMapping("timestamp with time zone"),
    datetime.timedelta: _TypeMapping("interval"),
    # Network types
    ipaddress.IPv4Network: _TypeMapping("cidr"),
    ipaddress.IPv6Network: _TypeMapping("cidr"),
    ipaddress.IPv4Address: _TypeMapping("inet"),
    ipaddress.IPv6Address: _TypeMapping("inet"),
    ipaddress.IPv4Interface: _TypeMapping("inet"),
    ipaddress.IPv6Interface: _TypeMapping("inet"),
}

# Default mapping for complex types that need JSON encoding
_JSONB_MAPPING = _TypeMapping("jsonb", _json_encoder)


async def _get_type_mapping(
    python_type: Any, *, vector_schema: res_schema.VectorSchema | None = None
) -> _TypeMapping:
    """
    Get the PostgreSQL type mapping for a Python type.

    Based on asyncpg's type conversion table:
    https://magicstack.github.io/asyncpg/current/usage.html#type-conversion

    For types that map to multiple PostgreSQL types, uses the broader one.
    Use `PgType` annotation with `typing.Annotated` to override the default.
    """
    type_info = analyze_type_info(python_type)

    # Check for PgType annotation override
    for annotation in type_info.annotations:
        if isinstance(annotation, PgType):
            return _TypeMapping(annotation.pg_type, annotation.encoder)

    base_type = type_info.base_type

    # Check direct leaf type mappings
    if base_type in _LEAF_TYPE_MAPPINGS:
        return _LEAF_TYPE_MAPPINGS[base_type]

    # NumPy ndarray: map to pgvector type bases; dimension is handled at the schema layer.
    if base_type is np.ndarray:
        if vector_schema is None:
            raise ValueError("VectorSpecProvider is required for NumPy ndarray type.")
        if vector_schema.size <= 0:
            raise ValueError(f"Invalid pgvector dimension: {vector_schema.size}")

        # Default to `vector` (float32/float64/int64/etc.). Use `halfvec` for float16.
        base = "halfvec" if vector_schema.dtype in (np.half, np.float16) else "vector"
        return _TypeMapping(pg_type=f"{base}({vector_schema.size})")

    elif vector_schema is not None:
        raise ValueError(
            f"VectorSpecProvider is only supported for NumPy ndarray type. Got type: {python_type}"
        )

    # Complex types that need JSON encoding
    if isinstance(
        type_info.variant, (SequenceType, MappingType, RecordType, UnionType, AnyType)
    ):
        return _JSONB_MAPPING

    # Default fallback
    return _JSONB_MAPPING


class ColumnDef(NamedTuple):
    """Definition of a table column."""

    type: str  # PostgreSQL type (e.g., "text", "bigint", "jsonb", "vector(384)")
    nullable: bool = True
    encoder: ValueEncoder | None = (
        None  # Optional encoder to convert value before sending to asyncpg
    )


# Type variable for row type
RowT = TypeVar("RowT", default=dict[str, Any])


@dataclass(slots=True)
class TableSchema(Generic[RowT]):
    """Schema definition for a PostgreSQL table."""

    columns: dict[str, ColumnDef]  # column name -> definition
    primary_key: list[str]  # Column names that form the primary key
    row_type: type[RowT] | None  # The row type, if provided

    def __init__(
        self,
        columns: dict[str, ColumnDef],
        primary_key: list[str],
        *,
        row_type: type[RowT] | None = None,
    ) -> None:
        """
        Create a TableSchema from pre-resolved column definitions.

        For constructing from a record type, use the async classmethod
        ``from_class`` instead.

        Args:
            columns: A dict mapping column names to ColumnDef.
            primary_key: List of column names that form the primary key.
            row_type: Optional original record type.
        """
        self.columns = columns
        self.primary_key = primary_key
        self.row_type = row_type

        # Validate primary key columns exist
        for pk in self.primary_key:
            if pk not in self.columns:
                raise ValueError(
                    f"Primary key column '{pk}' not found in columns: {list(self.columns.keys())}"
                )

    @classmethod
    async def from_class(
        cls,
        record_type: type[RowT],
        primary_key: list[str],
        *,
        column_overrides: dict[str, PgType | res_schema.VectorSchemaProvider]
        | None = None,
    ) -> "TableSchema[RowT]":
        """
        Create a TableSchema from a record type (dataclass, NamedTuple, or Pydantic model).

        Python types are automatically mapped to PostgreSQL types based on asyncpg's
        type conversion.

        Args:
            record_type: A record type (dataclass, NamedTuple, or Pydantic model).
            primary_key: List of column names that form the primary key.
            column_overrides: Optional dict mapping column names to PgType or
                              VectorSchemaProvider to override the default type mapping.
        """
        if not is_record_type(record_type):
            raise TypeError(
                f"record_type must be a record type (dataclass, NamedTuple, Pydantic model), "
                f"got {type(record_type)}"
            )
        columns = await cls._columns_from_record_type(record_type, column_overrides)
        return cls(columns, primary_key, row_type=record_type)

    @staticmethod
    async def _columns_from_record_type(
        record_type: type,
        column_overrides: dict[str, PgType | res_schema.VectorSchemaProvider] | None,
    ) -> dict[str, ColumnDef]:
        """Convert a record type to a dict of column name -> ColumnDef."""
        record_info = RecordType(record_type)
        columns: dict[str, ColumnDef] = {}

        for field in record_info.fields:
            type_info = analyze_type_info(field.type_hint)

            all_annotations = []
            if (
                override := column_overrides and column_overrides.get(field.name)
            ) is not None:
                all_annotations.append(override)
            all_annotations.extend(type_info.annotations)

            # Extract PgType and VectorSchemaProvider from annotations
            pg_type_annotation = next(
                (t for t in all_annotations if isinstance(t, PgType)), None
            )
            vector_schema = await anext(
                (
                    s
                    for annot in all_annotations
                    if (s := await res_schema.get_vector_schema(annot)) is not None
                ),
                None,
            )

            # Determine type mapping
            if pg_type_annotation is not None:
                type_mapping = _TypeMapping(
                    pg_type_annotation.pg_type, pg_type_annotation.encoder
                )
            else:
                type_mapping = await _get_type_mapping(
                    field.type_hint, vector_schema=vector_schema
                )

            columns[field.name] = ColumnDef(
                type=type_mapping.pg_type.strip(),
                nullable=type_info.nullable,
                encoder=type_mapping.encoder,
            )

        return columns


class _RowAction(NamedTuple):
    """Action to perform on a row."""

    key: _RowKey
    value: _RowValue | None  # None means delete


# --- Vector Index Attachment ---

_METRIC_OP_CLASS: dict[str, str] = {
    "cosine": "vector_cosine_ops",
    "l2": "vector_l2_ops",
    "ip": "vector_ip_ops",
}


class _VectorIndexSpec(NamedTuple):
    column: str
    metric: str
    method: str
    lists: int | None
    m: int | None
    ef_construction: int | None


_VectorIndexFingerprint = bytes


class _VectorIndexAction(NamedTuple):
    name: str
    spec: _VectorIndexSpec | None  # None means delete


class _VectorIndexHandler:
    """Handler for vector index attachment states."""

    _pool: asyncpg.Pool
    _table_name: str
    _schema_name: str | None
    _sink: coco.TargetActionSink[_VectorIndexAction]

    def __init__(
        self,
        pool: asyncpg.Pool,
        table_name: str,
        pg_schema_name: str | None,
    ) -> None:
        self._pool = pool
        self._table_name = table_name
        self._schema_name = pg_schema_name
        self._sink = coco.TargetActionSink.from_async_fn(self._apply_actions)

    async def _apply_actions(
        self, context_provider: ContextProvider, actions: Sequence[_VectorIndexAction]
    ) -> None:
        async with self._pool.acquire() as conn:
            for action in actions:
                index_name = f'"{self._table_name}__vector__{action.name}"'
                if action.spec is None:
                    await conn.execute(f"DROP INDEX IF EXISTS {index_name}")
                else:
                    # Drop + recreate
                    await conn.execute(f"DROP INDEX IF EXISTS {index_name}")
                    table_name = _qualified_table_name(
                        self._table_name, self._schema_name
                    )
                    op_class = _METRIC_OP_CLASS[action.spec.metric]
                    with_params: list[str] = []
                    if action.spec.method == "ivfflat":
                        if action.spec.lists is not None:
                            with_params.append(f"lists = {action.spec.lists}")
                    elif action.spec.method == "hnsw":
                        if action.spec.m is not None:
                            with_params.append(f"m = {action.spec.m}")
                        if action.spec.ef_construction is not None:
                            with_params.append(
                                f"ef_construction = {action.spec.ef_construction}"
                            )
                    with_clause = (
                        f" WITH ({', '.join(with_params)})" if with_params else ""
                    )
                    sql = (
                        f"CREATE INDEX {index_name} ON {table_name} "
                        f'USING {action.spec.method} ("{action.spec.column}" {op_class})'
                        f"{with_clause}"
                    )
                    await conn.execute(sql)

    def reconcile(
        self,
        key: coco.StableKey,
        desired_state: _VectorIndexSpec | coco.NonExistenceType,
        prev_possible_states: Collection[_VectorIndexFingerprint],
        prev_may_be_missing: bool,
        /,
    ) -> coco.TargetReconcileOutput[_VectorIndexAction, _VectorIndexFingerprint] | None:
        assert isinstance(key, str)
        if coco.is_non_existence(desired_state):
            if not prev_possible_states and not prev_may_be_missing:
                return None
            return coco.TargetReconcileOutput(
                action=_VectorIndexAction(name=key, spec=None),
                sink=self._sink,
                tracking_record=coco.NON_EXISTENCE,
            )

        target_fp = fingerprint_object(desired_state)
        if not prev_may_be_missing and all(
            prev == target_fp for prev in prev_possible_states
        ):
            return None

        return coco.TargetReconcileOutput(
            action=_VectorIndexAction(name=key, spec=desired_state),
            sink=self._sink,
            tracking_record=target_fp,
        )


# --- SQL Command Attachment ---


class _SqlCommandSpec(NamedTuple):
    setup_sql: str
    teardown_sql: str | None


class _SqlCommandAction(NamedTuple):
    name: str
    spec: _SqlCommandSpec | None  # None means delete
    prev_teardown_sql: str | None


def _collect_teardown_sql(
    prev_possible_states: Collection[_SqlCommandSpec],
) -> str | None:
    """Extract the first non-None teardown_sql from previous states."""
    for prev in prev_possible_states:
        if prev.teardown_sql is not None:
            return prev.teardown_sql
    return None


class _SqlCommandHandler:
    """Handler for SQL command attachment states."""

    _pool: asyncpg.Pool
    _table_name: str
    _schema_name: str | None
    _sink: coco.TargetActionSink[_SqlCommandAction]

    def __init__(
        self,
        pool: asyncpg.Pool,
        table_name: str,
        pg_schema_name: str | None,
    ) -> None:
        self._pool = pool
        self._table_name = table_name
        self._schema_name = pg_schema_name
        self._sink = coco.TargetActionSink.from_async_fn(self._apply_actions)

    async def _apply_actions(
        self, context_provider: ContextProvider, actions: Sequence[_SqlCommandAction]
    ) -> None:
        async with self._pool.acquire() as conn:
            for action in actions:
                # Run teardown of previous state if applicable
                if action.prev_teardown_sql is not None:
                    await conn.execute(action.prev_teardown_sql)
                # Run setup of new state if applicable
                if action.spec is not None:
                    await conn.execute(action.spec.setup_sql)

    def reconcile(
        self,
        key: coco.StableKey,
        desired_state: _SqlCommandSpec | coco.NonExistenceType,
        prev_possible_states: Collection[_SqlCommandSpec],
        prev_may_be_missing: bool,
        /,
    ) -> coco.TargetReconcileOutput[_SqlCommandAction, _SqlCommandSpec] | None:
        assert isinstance(key, str)
        if coco.is_non_existence(desired_state):
            if not prev_possible_states and not prev_may_be_missing:
                return None
            prev_teardown = _collect_teardown_sql(prev_possible_states)
            return coco.TargetReconcileOutput(
                action=_SqlCommandAction(
                    name=key, spec=None, prev_teardown_sql=prev_teardown
                ),
                sink=self._sink,
                tracking_record=coco.NON_EXISTENCE,
            )

        if not prev_may_be_missing and all(
            prev == desired_state for prev in prev_possible_states
        ):
            return None

        prev_teardown = _collect_teardown_sql(prev_possible_states)
        return coco.TargetReconcileOutput(
            action=_SqlCommandAction(
                name=key, spec=desired_state, prev_teardown_sql=prev_teardown
            ),
            sink=self._sink,
            tracking_record=desired_state,
        )


class _RowHandler(coco.TargetHandler[_RowValue, _RowFingerprint]):
    """Handler for row-level target states within a table."""

    _pool: asyncpg.Pool
    _table_name: str
    _schema_name: str | None
    _table_schema: TableSchema
    _sink: coco.TargetActionSink[_RowAction]

    def __init__(
        self,
        pool: asyncpg.Pool,
        table_name: str,
        pg_schema_name: str | None,
        table_schema: TableSchema,
    ) -> None:
        self._pool = pool
        self._table_name = table_name
        self._schema_name = pg_schema_name
        self._table_schema = table_schema
        self._sink = coco.TargetActionSink.from_async_fn(self._apply_actions)

    async def _apply_actions(
        self, context_provider: ContextProvider, actions: Sequence[_RowAction]
    ) -> None:
        """Apply row actions (upserts and deletes) to the database."""

        if not actions:
            return

        upserts: list[_RowAction] = []
        deletes: list[_RowAction] = []

        for action in actions:
            if action.value is None:
                deletes.append(action)
            else:
                upserts.append(action)

        async with asyncio.TaskGroup() as tg:
            self._schedule_upserts(tg, upserts)
            self._schedule_deletes(tg, deletes)

    def _schedule_upserts(
        self,
        tg: asyncio.TaskGroup,
        upserts: list[_RowAction],
    ) -> None:
        """Schedule upsert chunks as parallel tasks."""
        if not upserts:
            return

        table_name = _qualified_table_name(self._table_name, self._schema_name)
        columns = self._table_schema.columns
        pk_cols = self._table_schema.primary_key
        all_col_names = list(columns.keys())
        non_pk_cols = [c for c in all_col_names if c not in pk_cols]

        col_list = ", ".join(f'"{c}"' for c in all_col_names)
        pk_list = ", ".join(f'"{c}"' for c in pk_cols)

        if non_pk_cols:
            update_list = ", ".join(f'"{c}" = EXCLUDED."{c}"' for c in non_pk_cols)
            conflict_clause = f"ON CONFLICT ({pk_list}) DO UPDATE SET {update_list}"
        else:
            conflict_clause = f"ON CONFLICT ({pk_list}) DO NOTHING"

        num_parameters = len(all_col_names)
        if num_parameters == 0:
            return

        chunk_size = max(1, _BIND_LIMIT // num_parameters)
        for i in range(0, len(upserts), chunk_size):
            chunk = upserts[i : i + chunk_size]
            tg.create_task(
                self._execute_upsert_chunk(
                    table_name,
                    col_list,
                    conflict_clause,
                    all_col_names,
                    num_parameters,
                    chunk,
                )
            )

    async def _execute_upsert_chunk(
        self,
        table_name: str,
        col_list: str,
        conflict_clause: str,
        all_col_names: list[str],
        num_parameters: int,
        chunk: list[_RowAction],
    ) -> None:
        """Execute a single upsert chunk."""
        values_sql_parts: list[str] = []
        params: list[Any] = []
        for row_idx, action in enumerate(chunk):
            assert action.value is not None
            base = row_idx * num_parameters
            placeholders = ", ".join(f"${base + j + 1}" for j in range(num_parameters))
            values_sql_parts.append(f"({placeholders})")
            params.extend(action.value.get(col_name) for col_name in all_col_names)

        values_sql = ", ".join(values_sql_parts)
        sql = f"INSERT INTO {table_name} ({col_list}) VALUES {values_sql} {conflict_clause}"
        async with self._pool.acquire() as conn:
            await conn.execute(sql, *params)

    def _schedule_deletes(
        self,
        tg: asyncio.TaskGroup,
        deletes: list[_RowAction],
    ) -> None:
        """Schedule delete chunks as parallel tasks."""
        if not deletes:
            return

        table_name = _qualified_table_name(self._table_name, self._schema_name)
        pk_cols = self._table_schema.primary_key
        num_pk = len(pk_cols)

        chunk_size = max(1, _BIND_LIMIT // num_pk)
        for i in range(0, len(deletes), chunk_size):
            chunk = deletes[i : i + chunk_size]
            tg.create_task(
                self._execute_delete_chunk(table_name, pk_cols, num_pk, chunk)
            )

    async def _execute_delete_chunk(
        self,
        table_name: str,
        pk_cols: list[str],
        num_pk: int,
        chunk: list[_RowAction],
    ) -> None:
        """Execute a single batched delete chunk."""
        params: list[Any] = []
        if num_pk == 1:
            placeholders = ", ".join(f"${i + 1}" for i in range(len(chunk)))
            sql = f'DELETE FROM {table_name} WHERE "{pk_cols[0]}" IN ({placeholders})'
            params.extend(action.key[0] for action in chunk)
        else:
            or_parts: list[str] = []
            for row_idx, action in enumerate(chunk):
                base = row_idx * num_pk
                and_parts = [f'"{pk_cols[j]}" = ${base + j + 1}' for j in range(num_pk)]
                or_parts.append(f"({' AND '.join(and_parts)})")
                params.extend(action.key)
            sql = f"DELETE FROM {table_name} WHERE {' OR '.join(or_parts)}"

        async with self._pool.acquire() as conn:
            await conn.execute(sql, *params)

    def attachment(
        self, att_type: str
    ) -> _VectorIndexHandler | _SqlCommandHandler | None:
        if att_type == "vector_index":
            return _VectorIndexHandler(self._pool, self._table_name, self._schema_name)
        if att_type == "sql_command_attachment":
            return _SqlCommandHandler(self._pool, self._table_name, self._schema_name)
        return None

    def reconcile(
        self,
        key: coco.StableKey,
        desired_state: _RowValue | coco.NonExistenceType,
        prev_possible_states: Collection[_RowFingerprint],
        prev_may_be_missing: bool,
        /,
    ) -> coco.TargetReconcileOutput[_RowAction, _RowFingerprint] | None:
        key = _ROW_KEY_CHECKER.check(key)
        if coco.is_non_existence(desired_state):
            # Delete case - only if it might exist
            if not prev_possible_states and not prev_may_be_missing:
                return None
            return coco.TargetReconcileOutput(
                action=_RowAction(key=key, value=None),
                sink=self._sink,
                tracking_record=coco.NON_EXISTENCE,
            )

        # Upsert case
        target_fp = fingerprint_object(desired_state)
        if not prev_may_be_missing and all(
            prev == target_fp for prev in prev_possible_states
        ):
            # No change needed
            return None

        return coco.TargetReconcileOutput(
            action=_RowAction(key=key, value=desired_state),
            sink=self._sink,
            tracking_record=target_fp,
        )


class _TableKey(NamedTuple):
    """Key identifying a table: (database_key, pg_schema_name, table_name)."""

    db_key: str  # Stable key for the database
    pg_schema_name: str | None
    table_name: str


_TABLE_KEY_CHECKER = TypeChecker(tuple[str, str | None, str])


@dataclass
class _TableSpec:
    """Specification for a PostgreSQL table."""

    table_schema: TableSchema[Any]
    managed_by: Literal["system", "user"] = "system"


@unpickle_safe
class _PkColumnTrackingRecord(NamedTuple):
    """Primary-key column signature used for table-level main tracking record."""

    name: str
    type: str


@unpickle_safe
class _NonPkColumnTrackingRecord(NamedTuple):
    """Per-non-PK column tracking record used for incremental ALTER TABLE operations."""

    type: str
    nullable: bool


_EXT_PGVECTOR_SUBKEY: str = "ext:pgvector"
_COL_SUBKEY_PREFIX: str = "col:"


def _schema_uses_pgvector(schema: TableSchema[Any]) -> bool:
    return any(_is_pgvector_pg_type(c.type) for c in schema.columns.values())


def _col_subkey(col_name: str) -> str:
    return f"{_COL_SUBKEY_PREFIX}{col_name}"


_TableSubTrackingRecord = _NonPkColumnTrackingRecord | None


def _table_composite_tracking_record_from_spec(
    spec: _TableSpec,
) -> statediff.CompositeTrackingRecord[
    tuple[_PkColumnTrackingRecord, ...], str, _TableSubTrackingRecord
]:
    schema = spec.table_schema
    col_by_name = schema.columns
    pk_sig = tuple(
        _PkColumnTrackingRecord(name=pk, type=col_by_name[pk].type)
        for pk in schema.primary_key
    )
    sub: dict[str, _TableSubTrackingRecord] = {
        _col_subkey(col_name): _NonPkColumnTrackingRecord(
            type=col_def.type, nullable=col_def.nullable
        )
        for col_name, col_def in schema.columns.items()
        if col_name not in schema.primary_key
    }
    if _schema_uses_pgvector(schema):
        sub[_EXT_PGVECTOR_SUBKEY] = None
    return statediff.CompositeTrackingRecord(main=pk_sig, sub=sub)


_TableTrackingRecord = statediff.MutualTrackingRecord[
    statediff.CompositeTrackingRecord[
        tuple[_PkColumnTrackingRecord, ...], str, _TableSubTrackingRecord
    ]
]


class _TableAction(NamedTuple):
    """Action to perform on a table."""

    key: _TableKey
    spec: _TableSpec | coco.NonExistenceType
    main_action: statediff.DiffAction | None
    column_actions: dict[str, statediff.DiffAction]


class _TableHandler(coco.TargetHandler[_TableSpec, _TableTrackingRecord, _RowHandler]):
    """Handler for table-level target states."""

    _sink: coco.TargetActionSink[_TableAction, _RowHandler]

    def __init__(self) -> None:
        self._sink = coco.TargetActionSink.from_async_fn(self._apply_actions)

    async def _apply_actions(
        self, context_provider: ContextProvider, actions: Collection[_TableAction]
    ) -> list[coco.ChildTargetDef[_RowHandler] | None]:
        """Apply table actions (DDL) and return child row handlers."""
        actions_list = list(actions)
        outputs: list[coco.ChildTargetDef[_RowHandler] | None] = [None] * len(
            actions_list
        )

        # Group actions by table key so we can apply all DDL for the same table
        # within a single transaction/connection.
        by_key: dict[_TableKey, list[int]] = {}
        for i, action in enumerate(actions_list):
            by_key.setdefault(action.key, []).append(i)

        for key, idxs in by_key.items():
            pool = context_provider.get(key.db_key, asyncpg.Pool)
            async with pool.acquire() as conn:
                async with conn.transaction():
                    for i in idxs:
                        action = actions_list[i]
                        assert action.key == key

                        if action.main_action in ("replace", "delete"):
                            await self._drop_table(
                                conn, key.table_name, key.pg_schema_name
                            )

                        if coco.is_non_existence(action.spec):
                            outputs[i] = None
                            continue

                        spec = action.spec
                        outputs[i] = coco.ChildTargetDef(
                            handler=_RowHandler(
                                pool=pool,
                                table_name=key.table_name,
                                pg_schema_name=key.pg_schema_name,
                                table_schema=spec.table_schema,
                            )
                        )

                        if action.main_action in ("insert", "upsert", "replace"):
                            await self._create_table(
                                conn,
                                key,
                                spec.table_schema,
                                if_not_exists=(action.main_action == "upsert"),
                            )
                            continue

                        # No main change: reconcile non-PK columns incrementally.
                        if action.column_actions:
                            await self._apply_column_actions(
                                conn, key, spec.table_schema, action.column_actions
                            )

        return outputs

    async def _drop_table(
        self,
        conn: asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
        table_name: str,
        pg_schema_name: str | None,
    ) -> None:
        """Drop a table if it exists."""
        qualified_name = _qualified_table_name(table_name, pg_schema_name)
        await conn.execute(f"DROP TABLE IF EXISTS {qualified_name}")

    async def _ensure_pgvector_extension(
        self,
        conn: asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
        pg_schema_name: str | None,
    ) -> None:
        """
        Ensure the pgvector extension is installed.

        Postgres extensions are installed per-database but can live in a chosen
        schema; when `pg_schema_name` is provided we request installing into it.
        """
        schema_clause = f' WITH SCHEMA "{pg_schema_name}"' if pg_schema_name else ""
        await conn.execute(f"CREATE EXTENSION IF NOT EXISTS vector{schema_clause}")

    async def _create_table(
        self,
        conn: asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
        key: _TableKey,
        schema: TableSchema[Any],
        *,
        if_not_exists: bool,
    ) -> None:
        """Create a table."""
        qualified_name = _qualified_table_name(key.table_name, key.pg_schema_name)

        # Create schema if specified
        if key.pg_schema_name:
            await conn.execute(f'CREATE SCHEMA IF NOT EXISTS "{key.pg_schema_name}"')

        # Ensure pgvector extension exists if needed by any column type.
        if _schema_uses_pgvector(schema):
            await self._ensure_pgvector_extension(conn, key.pg_schema_name)

        # Build column definitions
        col_defs = []
        for col_name, col in schema.columns.items():
            nullable = (
                ""
                if col.nullable and col_name not in schema.primary_key
                else " NOT NULL"
            )
            col_defs.append(f'"{col_name}" {col.type}{nullable}')

        # Build primary key constraint
        pk_cols = ", ".join(f'"{c}"' for c in schema.primary_key)
        col_defs.append(f"PRIMARY KEY ({pk_cols})")

        columns_sql = ", ".join(col_defs)
        if_not_exists_sql = " IF NOT EXISTS" if if_not_exists else ""
        sql = f"CREATE TABLE{if_not_exists_sql} {qualified_name} ({columns_sql})"
        await conn.execute(sql)

    async def _apply_column_actions(
        self,
        conn: asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
        key: _TableKey,
        schema: TableSchema[Any],
        column_actions: dict[str, statediff.DiffAction],
    ) -> None:
        qualified_name = _qualified_table_name(key.table_name, key.pg_schema_name)
        pk_cols = set(schema.primary_key)
        non_pk_col_by_name = {
            n: c for n, c in schema.columns.items() if n not in pk_cols
        }
        for sub_key, action in column_actions.items():
            if sub_key == _EXT_PGVECTOR_SUBKEY:
                if action != "delete":
                    await self._ensure_pgvector_extension(conn, key.pg_schema_name)
                continue

            if not sub_key.startswith(_COL_SUBKEY_PREFIX):
                raise ValueError(
                    f"Unexpected column subkey format: {sub_key!r}, expected to start with {_COL_SUBKEY_PREFIX!r}"
                )
            col_name = sub_key[len(_COL_SUBKEY_PREFIX) :]

            # Defensive: we never ALTER PK columns here.
            if col_name in pk_cols:
                continue

            if action == "delete":
                await conn.execute(
                    f'ALTER TABLE {qualified_name} DROP COLUMN IF EXISTS "{col_name}"'
                )
                continue

            desired_col = non_pk_col_by_name.get(col_name)
            if desired_col is None:
                # If the desired schema no longer mentions this column, treat
                # it as a no-op here; "delete" should have been emitted.
                continue

            if action == "insert":
                async with conn.transaction():
                    await conn.execute(
                        f"ALTER TABLE {qualified_name} "
                        f'ADD COLUMN "{col_name}" {desired_col.type}'
                    )
                continue

            if action == "upsert":
                await conn.execute(
                    f"ALTER TABLE {qualified_name} "
                    f'ADD COLUMN IF NOT EXISTS "{col_name}" {desired_col.type}'
                )
                continue

            if action == "replace":
                # Type change may fail depending on existing data. Try ALTER TYPE first
                # inside a savepoint; if it fails, fall back to drop+add.
                try:
                    async with conn.transaction():
                        nullable = "" if desired_col.nullable else " NOT NULL"
                        await conn.execute(
                            f"ALTER TABLE {qualified_name} "
                            f'ALTER COLUMN "{col_name}" TYPE {desired_col.type}{nullable}'
                        )
                except asyncpg.PostgresError:
                    await conn.execute(
                        f'ALTER TABLE {qualified_name} DROP COLUMN IF EXISTS "{col_name}"'
                    )
                    await conn.execute(
                        f"ALTER TABLE {qualified_name} "
                        f'ADD COLUMN "{col_name}" {desired_col.type}'
                    )

    def reconcile(
        self,
        key: coco.StableKey,
        desired_state: _TableSpec | coco.NonExistenceType,
        prev_possible_states: Collection[_TableTrackingRecord],
        prev_may_be_missing: bool,
        /,
    ) -> (
        coco.TargetReconcileOutput[_TableAction, _TableTrackingRecord, _RowHandler]
        | None
    ):
        key = _TableKey(*_TABLE_KEY_CHECKER.check(key))

        tracking_record: _TableTrackingRecord | coco.NonExistenceType

        if coco.is_non_existence(desired_state):
            tracking_record = coco.NON_EXISTENCE
        else:
            tracking_record = statediff.MutualTrackingRecord(
                tracking_record=_table_composite_tracking_record_from_spec(
                    desired_state
                ),
                managed_by=desired_state.managed_by,
            )

        resolved = statediff.resolve_system_transition(
            statediff.TrackingRecordTransition(
                tracking_record,
                prev_possible_states,
                prev_may_be_missing,
            )
        )
        main_action, column_transitions = statediff.diff_composite(resolved)

        column_actions: dict[str, statediff.DiffAction] = {}
        if main_action is None:
            for sub_key, t in column_transitions.items():
                action = statediff.diff(t)
                if action is not None:
                    column_actions[sub_key] = action

        # Determine child invalidation for row-level targets.
        child_invalidation: Literal["destructive", "lossy"] | None = None
        if main_action == "replace":
            # Table is dropped and recreated — all rows are destroyed.
            child_invalidation = "destructive"
        elif main_action is None and any(
            a != "insert" for a in column_actions.values()
        ):
            # Column schema changes (other than adding new columns) may lose data.
            child_invalidation = "lossy"

        return coco.TargetReconcileOutput(
            action=_TableAction(
                key=key,
                spec=desired_state,
                main_action=main_action,
                column_actions=column_actions,
            ),
            sink=self._sink,
            tracking_record=tracking_record,
            child_invalidation=child_invalidation,
        )


# Register the root target states provider
_table_provider = coco.register_root_target_states_provider(
    "cocoindex/postgres/table", _TableHandler()
)


class TableTarget(
    Generic[RowT, coco.MaybePendingS], coco.ResolvesTo["TableTarget[RowT]"]
):
    """
    A target for writing rows to a PostgreSQL table.

    The table is managed as a target state, with the scope used to scope the target state.

    Type Parameters:
        RowT: The type of row objects (dict, dataclass, NamedTuple, or Pydantic model).
    """

    _provider: coco.TargetStateProvider[_RowValue, None, coco.MaybePendingS]
    _table_schema: TableSchema[RowT]

    def __init__(
        self,
        provider: coco.TargetStateProvider[_RowValue, None, coco.MaybePendingS],
        table_schema: TableSchema[RowT],
    ) -> None:
        self._provider = provider
        self._table_schema = table_schema

    def declare_row(self: "TableTarget[RowT]", *, row: RowT) -> None:
        """
        Declare a row to be upserted to this table.

        Args:
            row: A row object (dict, dataclass, NamedTuple, or Pydantic model).
                 Must include all primary key columns.
        """
        row_dict = self._row_to_dict(row)
        # Extract primary key values
        pk_values = tuple(row_dict[pk] for pk in self._table_schema.primary_key)
        coco.declare_target_state(self._provider.target_state(pk_values, row_dict))

    def _row_to_dict(self, row: RowT) -> dict[str, Any]:
        """
        Convert a row (dict or object) into dict[str, Any] using the schema columns,
        and apply column encoders for both dict and object inputs.
        """
        out: dict[str, Any] = {}
        for col_name, col in self._table_schema.columns.items():
            if isinstance(row, dict):
                value = row.get(col_name)
            else:
                value = getattr(row, col_name)

            if value is not None and col.encoder is not None:
                value = col.encoder(value)
            out[col_name] = value
        return out

    def declare_vector_index(
        self: "TableTarget[RowT]",
        *,
        name: str | None = None,
        column: str,
        metric: Literal["cosine", "l2", "ip"] = "cosine",
        method: Literal["ivfflat", "hnsw"] = "ivfflat",
        lists: int | None = None,
        m: int | None = None,
        ef_construction: int | None = None,
    ) -> None:
        """
        Declare a pgvector index on a column of this table.

        The actual Postgres index will be named ``{table_name}__vector__{name}``.

        Args:
            name: Logical index name (defaults to ``column``).
            column: Column to index.
            metric: Distance metric ("cosine", "l2", or "ip").
            method: Index method ("ivfflat" or "hnsw").
            lists: Number of lists (ivfflat only).
            m: Maximum number of connections per layer (hnsw only).
            ef_construction: Size of the dynamic candidate list (hnsw only).
        """
        if name is None:
            name = column
        spec = _VectorIndexSpec(
            column=column,
            metric=metric,
            method=method,
            lists=lists,
            m=m,
            ef_construction=ef_construction,
        )
        att_provider = self._provider.attachment("vector_index")
        coco.declare_target_state(att_provider.target_state(name, spec))

    def declare_sql_command_attachment(
        self: "TableTarget[RowT]",
        *,
        name: str,
        setup_sql: str,
        teardown_sql: str | None = None,
    ) -> None:
        """
        Declare a SQL command attachment on this table.

        The setup SQL is executed when the attachment is created or changed.
        The teardown SQL (if provided) is executed when the attachment is removed
        or before re-running setup on change.

        Args:
            name: Attachment name (stable identifier).
            setup_sql: SQL to execute on creation/change.
            teardown_sql: SQL to execute on removal/before change (optional).
        """
        spec = _SqlCommandSpec(
            setup_sql=setup_sql,
            teardown_sql=teardown_sql,
        )
        att_provider = self._provider.attachment("sql_command_attachment")
        coco.declare_target_state(att_provider.target_state(name, spec))

    def __coco_memo_key__(self) -> str:
        return self._provider.memo_key


def table_target(
    db: ContextKey[asyncpg.Pool],
    table_name: str,
    table_schema: TableSchema[RowT],
    *,
    pg_schema_name: str | None = None,
    managed_by: Literal["system", "user"] = "system",
) -> coco.TargetState[_RowHandler]:
    """
    Create a TargetState for a PostgreSQL table target.

    Use with ``coco.mount_target()`` to mount and get a child provider,
    or with ``mount_table_target()`` for a convenience wrapper.

    Args:
        db: ContextKey for the asyncpg.Pool connection.
        table_name: Name of the table.
        table_schema: Schema definition including columns and primary key.
        pg_schema_name: Optional PostgreSQL schema name (default is "public").
        managed_by: Whether the table is managed by "system" or "user".

    Returns:
        A TargetState that can be passed to ``mount_target()``.
    """
    key = _TableKey(
        db_key=db.key,
        pg_schema_name=pg_schema_name,
        table_name=table_name,
    )
    spec = _TableSpec(
        table_schema=table_schema,
        managed_by=managed_by,
    )
    return _table_provider.target_state(key, spec)


def declare_table_target(
    db: ContextKey[asyncpg.Pool],
    table_name: str,
    table_schema: TableSchema[RowT],
    *,
    pg_schema_name: str | None = None,
    managed_by: Literal["system", "user"] = "system",
) -> TableTarget[RowT, coco.PendingS]:
    """
    Create a TableTarget for writing rows to a PostgreSQL table.

    Args:
        db: ContextKey for the asyncpg.Pool connection.
        table_name: Name of the table.
        table_schema: Schema definition including columns and primary key.
        pg_schema_name: Optional PostgreSQL schema name (default is "public").
        managed_by: Whether the table is managed by "system" (CocoIndex creates/drops it)
                    or "user" (table must exist, CocoIndex only manages rows).

    Returns:
        A TableTarget that can be used to declare rows.
    """
    provider = coco.declare_target_state_with_child(
        table_target(
            db,
            table_name,
            table_schema,
            pg_schema_name=pg_schema_name,
            managed_by=managed_by,
        )
    )
    return TableTarget(provider, table_schema)


async def mount_table_target(
    db: ContextKey[asyncpg.Pool],
    table_name: str,
    table_schema: TableSchema[RowT],
    *,
    pg_schema_name: str | None = None,
    managed_by: Literal["system", "user"] = "system",
) -> TableTarget[RowT]:
    """
    Mount a table target and return a ready-to-use TableTarget.

    Sugar over ``table_target()`` + ``coco.mount_target()`` + wrapping.

    Args:
        db: ContextKey for the asyncpg.Pool connection.
        table_name: Name of the table.
        table_schema: Schema definition including columns and primary key.
        pg_schema_name: Optional PostgreSQL schema name (default is "public").
        managed_by: Whether the table is managed by "system" or "user".

    Returns:
        A TableTarget that can be used to declare rows.
    """
    provider = await coco.mount_target(
        table_target(
            db,
            table_name,
            table_schema,
            pg_schema_name=pg_schema_name,
            managed_by=managed_by,
        )
    )
    return TableTarget(provider, table_schema)


async def create_pool(
    dsn: str | None = None,
    *,
    init: Callable[[asyncpg.Connection], Any] | None = None,
    **kwargs: Any,
) -> asyncpg.Pool:
    """
    Create an asyncpg connection pool, registering extensions needed by the postgres connector on each connection.

    Args:
        dsn: Connection string or None.
        init: Optional connection initialization callback. If provided,
              it will be called for each connection in the pool.
        **kwargs: All other arguments are passed to `asyncpg.create_pool()`.

    Returns:
        An asyncpg connection pool.

    Example:
        ```python
        pool = await create_pool("postgresql://localhost/mydb")
        ```
    """

    async def _init_with_pgvector(conn: asyncpg.Connection) -> None:
        await pgvector.asyncpg.register_vector(conn)
        if init is not None:
            result = init(conn)
            if inspect.isawaitable(result):
                await result

    return await asyncpg.create_pool(dsn, init=_init_with_pgvector, **kwargs)


__all__ = [
    "ColumnDef",
    "ValueEncoder",
    "PgType",
    "TableSchema",
    "TableTarget",
    "create_pool",
    "declare_table_target",
    "mount_table_target",
    "table_target",
]
