from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Sequence
from typing import ClassVar

import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.schema import CreateSchema, CreateTable, DropSchema

from tracecat.auth.types import Role
from tracecat.custom_fields.schemas import CustomFieldCreate, CustomFieldUpdate
from tracecat.db.locks import derive_lock_key_from_parts, pg_advisory_lock
from tracecat.exceptions import TracecatValidationError
from tracecat.identifiers.workflow import WorkspaceUUID
from tracecat.service import BaseWorkspaceService
from tracecat.tables.service import (
    DYNAMIC_WORKSPACE_RLS_POLICY,
    DYNAMIC_WORKSPACE_TENANT_COLUMN,
    RLS_BYPASS_ON,
    RLS_BYPASS_VAR,
    RLS_WORKSPACE_VAR,
    TableEditorService,
    is_internal_column_name,
    sanitize_identifier,
    validate_identifier,
)


class CustomFieldsService(BaseWorkspaceService, ABC):
    """Base service for managing workspace-specific custom fields.

    Subclasses must define:
        - service_name: Identifier for the service
        - _table: The SQLAlchemy table name
        - _schema_prefix: Prefix for the workspace schema
        - _reserved_columns: Set of column names that cannot be deleted
        - initialize_workspace_schema(): Create the schema and base table
    """

    service_name: ClassVar[str]
    _table: ClassVar[str]
    _schema_prefix: ClassVar[str] = "custom_fields_"
    _reserved_columns: ClassVar[set[str]]

    def __init__(self, session: AsyncSession, role: Role | None = None):
        super().__init__(session, role)
        self._workspace_uuid = WorkspaceUUID.new(self.workspace_id)
        self.schema_name = self._get_schema_name()
        self._schema_initialized = False
        self.sanitized_table_name = sanitize_identifier(self._table)
        self.editor = TableEditorService(
            self.session,
            self.role,
            table_name=self._table,
            schema_name=self.schema_name,
        )

    def _get_schema_name(self) -> str:
        """Generate the schema name for this workspace."""

        return f"{self._schema_prefix}{self._workspace_uuid.short()}"

    def _full_table_name(self) -> str:
        """Get the fully-qualified workspace table name."""
        return f'"{self.schema_name}".{self.sanitized_table_name}'

    def _workspace_tenant_default_sql(self) -> sa.TextClause:
        """Build the server default expression for the tenant column."""
        return sa.text(f"'{self._workspace_uuid}'::uuid")

    async def _enable_workspace_rls_for_base_table(self) -> None:
        """Enable workspace RLS policy on the workspace base table."""
        conn = await self.session.connection()
        policy_expr = (
            f"current_setting('{RLS_BYPASS_VAR}', true) = '{RLS_BYPASS_ON}' "
            f'OR "{DYNAMIC_WORKSPACE_TENANT_COLUMN}" = '
            f"NULLIF(current_setting('{RLS_WORKSPACE_VAR}', true), '')::uuid"
        )
        full_table_name = self._full_table_name()
        await conn.execute(
            sa.DDL("ALTER TABLE %s ENABLE ROW LEVEL SECURITY", full_table_name)
        )
        await conn.execute(
            sa.DDL(
                f"DROP POLICY IF EXISTS {DYNAMIC_WORKSPACE_RLS_POLICY} ON %s",
                full_table_name,
            )
        )
        await conn.execute(
            sa.DDL(
                f"""
                CREATE POLICY {DYNAMIC_WORKSPACE_RLS_POLICY} ON %s
                    FOR ALL
                    USING ({policy_expr})
                    WITH CHECK ({policy_expr})
                """,
                full_table_name,
            )
        )

    def _assert_user_field_name_allowed(self, field_name: str) -> None:
        """Reject operations on internal/system-managed field names."""
        if is_internal_column_name(field_name):
            raise ValueError(f"Field {field_name} is reserved for internal use")

    @abstractmethod
    def _table_definition(self) -> sa.Table:
        """Return the SQLAlchemy Table definition for the workspace table.

        Subclasses must implement this to define the table structure.
        """

        raise NotImplementedError

    async def initialize_workspace_schema(self) -> None:
        """Create the workspace schema and base table if absent."""

        lock_key = derive_lock_key_from_parts(
            "tracecat.custom_fields.initialize_workspace_schema",
            self.service_name,
            self._table,
        )
        async with pg_advisory_lock(self.session, lock_key):
            await self.session.execute(
                CreateSchema(self.schema_name, if_not_exists=True)
            )
            await self.session.execute(
                CreateTable(self._table_definition(), if_not_exists=True)
            )
            await self._enable_workspace_rls_for_base_table()
        self._schema_initialized = True

    async def _ensure_schema_ready(self) -> None:
        """Ensure the workspace schema exists, creating it if needed."""

        if self._schema_initialized:
            return

        conn = await self.session.connection()

        def check_schema_exists(sync_conn: sa.Connection) -> bool:
            inspector = sa.inspect(sync_conn)
            return inspector.has_schema(self.schema_name) and inspector.has_table(
                self._table, schema=self.schema_name
            )

        if await conn.run_sync(check_schema_exists):
            self._schema_initialized = True
            return

        await self.initialize_workspace_schema()

    async def drop_workspace_schema(self) -> None:
        """Drop the workspace schema and all contained objects."""

        await self.session.execute(
            DropSchema(self.schema_name, cascade=True, if_exists=True)
        )
        self._schema_initialized = False

    async def list_fields(self) -> Sequence[sa.engine.interfaces.ReflectedColumn]:
        """List all custom fields for the workspace."""

        await self._ensure_schema_ready()
        return await self.editor.get_columns()

    async def create_field(self, params: CustomFieldCreate) -> None:
        """Create a new custom field column."""

        await self._ensure_schema_ready()
        params.nullable = True  # Custom fields remain nullable by default
        await self.editor.create_column(params)
        await self.session.commit()

    async def update_field(self, field_id: str, params: CustomFieldUpdate) -> None:
        """Update a custom field column."""

        await self._ensure_schema_ready()
        await self.editor.update_column(field_id, params)
        await self.session.commit()

    async def delete_field(self, field_id: str) -> None:
        """Delete a custom field.

        Args:
            field_id: The name of the field to delete

        Raises:
            TracecatValidationError: If the field is reserved or internal
        """

        await self._ensure_schema_ready()
        self._assert_user_field_name_allowed(field_id)
        if field_id in self._reserved_columns:
            raise TracecatValidationError(f"Field {field_id} is a reserved field")
        validate_identifier(field_id)
        await self.editor.delete_column(field_id)
        await self.session.commit()


__all__ = ["CustomFieldsService"]
