"""PostgreSQL schema and initialization for the bitable subsystem.

Uses an independent ``bitable`` schema within the shared PostgreSQL instance.
Follows the lazy-init + lock pattern from ``evolution/pg_store.py`` and the
``_SCHEMA_VERSION`` migration pattern from ``server/auth/models.py``.

Schema versioning
-----------------
:data:`_SCHEMA_VERSION` tracks the current bitable DB schema. The
``bitable_meta`` table stores the version so subsequent restarts skip
already-applied migrations.
"""

from __future__ import annotations

import asyncio
import logging
import os
import uuid as _uuid
from datetime import datetime, timezone
from types import TracebackType

from sqlalchemy import (
    Column,
    DateTime,
    Index,
    String,
    Text,
    UniqueConstraint,
    text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import DeclarativeBase

logger = logging.getLogger(__name__)

# Current schema version — bump when adding migrations.
# V1: initial schema (tables/fields/records/views/recalc/meta).
# V2: add bitable_files table + file_id column on bitable_tables (R1).
_SCHEMA_VERSION = 2

_META_SCHEMA_VERSION_KEY = "schema_version"


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


def _uuid_str() -> str:
    return str(_uuid.uuid4())


class BitableBase(DeclarativeBase):
    """Declarative base for bitable ORM models (independent schema)."""


class FileModel(BitableBase):
    """ORM model for ``bitable.bitable_files`` — top-level container (R1)."""

    __tablename__ = "bitable_files"
    __table_args__ = (
        Index("ix_bitable_files_owner", "owner_user_id"),
        {"schema": "bitable"},
    )

    id = Column(String, primary_key=True, default=_uuid_str)
    name = Column(String, nullable=False)
    icon = Column(String, default="table")
    description = Column(Text, default="")
    owner_user_id = Column(String, nullable=True)
    created_at = Column(DateTime(timezone=True), default=_utcnow)
    updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow)


class TableModel(BitableBase):
    """ORM model for ``bitable.bitable_tables``."""

    __tablename__ = "bitable_tables"
    __table_args__ = (
        Index("ix_bitable_tables_file_id", "file_id"),
        {"schema": "bitable"},
    )

    id = Column(String, primary_key=True, default=_uuid_str)
    file_id = Column(String, nullable=True)
    name = Column(String, nullable=False)
    description = Column(Text, default="")
    primary_key_field_id = Column(String, nullable=True)
    owner_user_id = Column(String, nullable=True)
    created_at = Column(DateTime(timezone=True), default=_utcnow)
    updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow)


class FieldModel(BitableBase):
    """ORM model for ``bitable.bitable_fields``."""

    __tablename__ = "bitable_fields"
    __table_args__ = (
        Index("ix_bitable_fields_table_id", "table_id"),
        {"schema": "bitable"},
    )

    id = Column(String, primary_key=True, default=_uuid_str)
    table_id = Column(String, nullable=False)
    name = Column(String, nullable=False)
    field_type = Column(String, nullable=False)
    config = Column(JSONB, default=dict)
    owner = Column(String, default="user")
    created_at = Column(DateTime(timezone=True), default=_utcnow)


class RecordModel(BitableBase):
    """ORM model for ``bitable.bitable_records``.

    ``values`` is JSONB mapping ``{field_id: value}``. A GIN index supports
    efficient JSONB key existence checks. A unique expression index on
    ``(table_id, values->>pk_field_id)`` enforces primary-key uniqueness
    (created dynamically in ``_apply_v1_schema`` because the pk field id is
    per-table, not a fixed column).
    """

    __tablename__ = "bitable_records"
    __table_args__ = (
        Index("ix_bitable_records_table_id", "table_id"),
        Index("ix_bitable_records_values_gin", "values", postgresql_using="gin"),
        {"schema": "bitable"},
    )

    id = Column(String, primary_key=True, default=_uuid_str)
    table_id = Column(String, nullable=False)
    values = Column(JSONB, default=dict)
    created_at = Column(DateTime(timezone=True), default=_utcnow)
    updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow)


class ViewModel(BitableBase):
    """ORM model for ``bitable.bitable_views``."""

    __tablename__ = "bitable_views"
    __table_args__ = (
        Index("ix_bitable_views_table_id", "table_id"),
        {"schema": "bitable"},
    )

    id = Column(String, primary_key=True, default=_uuid_str)
    table_id = Column(String, nullable=False)
    name = Column(String, nullable=False)
    view_type = Column(String, default="grid")
    config = Column(JSONB, default=dict)
    created_at = Column(DateTime(timezone=True), default=_utcnow)


class RecalcQueueModel(BitableBase):
    """ORM model for ``bitable.bitable_recalc_queue``.

    The ``(record_id, field_id)`` unique index prevents duplicate enqueues.
    The ``(status, queued_at)`` index supports efficient worker consumption.
    """

    __tablename__ = "bitable_recalc_queue"
    __table_args__ = (
        UniqueConstraint("record_id", "field_id", name="uq_recalc_record_field"),
        Index("ix_recalc_status_queued", "status", "queued_at"),
        {"schema": "bitable"},
    )

    id = Column(String, primary_key=True, default=_uuid_str)
    table_id = Column(String, nullable=False)
    record_id = Column(String, nullable=False)
    field_id = Column(String, nullable=False)
    status = Column(String, default="pending")
    error_message = Column(Text, nullable=True)
    queued_at = Column(DateTime(timezone=True), default=_utcnow)
    completed_at = Column(DateTime(timezone=True), nullable=True)


class MetaModel(BitableBase):
    """ORM model for ``bitable.bitable_meta`` — schema version tracking."""

    __tablename__ = "bitable_meta"
    __table_args__ = {"schema": "bitable"}

    key = Column(String, primary_key=True)
    value = Column(String, nullable=False)
    updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow)


# ---------------------------------------------------------------------------
# Schema is created via BitableBase.metadata.create_all (see BitableDB.init).
# The ORM models above are the single source of truth for table/index DDL.
# ---------------------------------------------------------------------------


async def _apply_v2_migration(conn: object) -> None:
    """V2 migration: create ``bitable_files`` table + add ``file_id`` to tables.

    Idempotent — safe to call on fresh installs (``create_all`` already made
    the files table and file_id column) and on V1 upgrades (ALTERs the
    existing tables). The CREATE TABLE IF NOT EXISTS + ADD COLUMN IF NOT
    EXISTS guards make both paths work.

    Defensive: V1 tables without a ``file_id`` are left with ``file_id = NULL``
    (orphaned). The service layer treats NULL file_id as "no parent file"
    and continues to work — they're reachable via the deprecated
    ``/tables`` endpoint. The plan notes there is no existing production
    data to migrate (verifier-confirmed), so this is purely defensive.
    """
    # bitable_files table — created by create_all on fresh installs, but
    # we recreate via raw SQL so V1→V2 upgrades (no create_all re-run for
    # existing tables) also get it.
    await conn.execute(
        text(
            "CREATE TABLE IF NOT EXISTS bitable.bitable_files ("
            "  id VARCHAR PRIMARY KEY,"
            "  name VARCHAR NOT NULL,"
            "  icon VARCHAR DEFAULT 'table',"
            "  description TEXT DEFAULT '',"
            "  owner_user_id VARCHAR,"
            "  created_at TIMESTAMPTZ DEFAULT NOW(),"
            "  updated_at TIMESTAMPTZ DEFAULT NOW()"
            ")"
        )
    )
    await conn.execute(
        text(
            "CREATE INDEX IF NOT EXISTS ix_bitable_files_owner "
            "ON bitable.bitable_files (owner_user_id)"
        )
    )

    # Add file_id column to bitable_tables (idempotent).
    # ponytail: ADD COLUMN IF NOT EXISTS is PostgreSQL 9.6+ — safe here.
    await conn.execute(
        text("ALTER TABLE bitable.bitable_tables ADD COLUMN IF NOT EXISTS file_id VARCHAR")
    )
    await conn.execute(
        text(
            "CREATE INDEX IF NOT EXISTS ix_bitable_tables_file_id "
            "ON bitable.bitable_tables (file_id)"
        )
    )
    logger.info("applied V2 migration: bitable_files table + file_id column")


def _resolve_database_url(database_url: str | None = None) -> str | None:
    """Resolve PostgreSQL connection URL.

    Priority: explicit arg > env ``DATABASE_URL`` > ``AGENTKIT_DATABASE_URL``.
    """
    if database_url:
        return database_url
    return os.environ.get("DATABASE_URL") or os.environ.get("AGENTKIT_DATABASE_URL")


class BitableDB:
    """PostgreSQL connection manager for bitable (lazy init + lock).

    Usage::

        db = BitableDB(database_url="postgresql+asyncpg://...")
        await db.init()
        # ... use db.engine / db.session_factory ...
        await db.close()
    """

    def __init__(self, database_url: str | None = None) -> None:
        self._database_url = database_url or _resolve_database_url()
        self._engine: object | None = None
        self._session_factory: object | None = None
        self._initialized = False
        self._init_lock = asyncio.Lock()

    @property
    def database_url(self) -> str | None:
        return self._database_url

    @property
    def engine(self) -> object | None:
        return self._engine

    @property
    def session_factory(self) -> object | None:
        return self._session_factory

    @property
    def is_initialized(self) -> bool:
        return self._initialized

    async def _ensure_initialized(self) -> None:
        """Lazy-init async engine and session factory (with lock)."""
        if self._initialized:
            return

        async with self._init_lock:
            if self._initialized:
                return

            if not self._database_url:
                raise RuntimeError(
                    "No database URL configured for bitable. "
                    "Set DATABASE_URL or AGENTKIT_DATABASE_URL env var."
                )

            from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
            from sqlalchemy.orm import sessionmaker

            self._engine = create_async_engine(self._database_url, echo=False)
            self._session_factory = sessionmaker(
                self._engine, class_=AsyncSession, expire_on_commit=False
            )
            self._initialized = True

    async def init(self) -> None:
        """Initialize engine + create schema and tables (idempotent).

        Runs schema migrations based on ``_SCHEMA_VERSION`` stored in
        ``bitable_meta``. Safe to call on every startup.
        """
        await self._ensure_initialized()

        async with self._engine.begin() as conn:
            # 1. Create the bitable schema (idempotent) — must precede
            #    metadata.create_all since tables live in this schema.
            await conn.execute(text("CREATE SCHEMA IF NOT EXISTS bitable"))

            # 2. Create all tables/indexes/constraints from ORM metadata.
            #    Uses run_sync because asyncpg doesn't support multi-statement
            #    text() execution; metadata.create_all emits one DDL statement
            #    per table and handles schema-qualified names + GIN indexes.
            await conn.run_sync(BitableBase.metadata.create_all)

            # 3. Check current schema version
            result = await conn.execute(
                text("SELECT value FROM bitable.bitable_meta WHERE key = :key"),
                {"key": _META_SCHEMA_VERSION_KEY},
            )
            row = result.fetchone()
            current_version = int(row[0]) if row else 0

            # 4. Apply migrations if needed (future versions add elif blocks here)
            if current_version < _SCHEMA_VERSION:
                # V1: initial schema (already created above via create_all)
                # V2: file layer (bitable_files table + file_id column on tables)
                if current_version < 2:
                    await _apply_v2_migration(conn)
                await conn.execute(
                    text(
                        "INSERT INTO bitable.bitable_meta (key, value, updated_at) "
                        "VALUES (:key, :value, NOW()) "
                        "ON CONFLICT (key) DO UPDATE SET value = :value, updated_at = NOW()"
                    ),
                    {"key": _META_SCHEMA_VERSION_KEY, "value": str(_SCHEMA_VERSION)},
                )
                logger.info("bitable schema migrated: v%d → v%d", current_version, _SCHEMA_VERSION)

    async def close(self) -> None:
        """Close engine, release all connections."""
        if self._engine is not None:
            await self._engine.dispose()
            self._engine = None
            self._session_factory = None
            self._initialized = False

    async def __aenter__(self) -> "BitableDB":
        await self.init()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        await self.close()


# Module-level singleton (initialized by app lifespan)
_db: BitableDB | None = None


async def init_bitable_db(database_url: str | None = None) -> BitableDB:
    """Initialize the bitable database (module-level singleton).

    Called from ``app.py`` lifespan. On failure, raises — the caller should
    catch and degrade gracefully (bitable API returns 503).
    """
    global _db
    if _db is not None and _db.is_initialized:
        return _db

    _db = BitableDB(database_url=database_url)
    await _db.init()
    logger.info("bitable DB initialized (schema version %d)", _SCHEMA_VERSION)
    return _db


def get_bitable_db() -> BitableDB | None:
    """Return the module-level bitable DB singleton (or None if not init'd)."""
    return _db


async def close_bitable_db() -> None:
    """Close the module-level bitable DB singleton."""
    global _db
    if _db is not None:
        await _db.close()
        _db = None