"""JSON query builder for KameleonDB.

Provides CRUD operations using JSON storage (JSONB on PostgreSQL, JSON on SQLite).
All field data is stored in a single JSON column on kdb_records.

Note: Complex queries should be generated by agents using get_schema_context()
and executed via execute_query() rather than programmatic filtering.
"""

from __future__ import annotations

from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
from uuid import uuid4

from sqlalchemy import cast, update
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Session

from kameleondb.exceptions import FieldNotFoundError, QueryError, RecordNotFoundError
from kameleondb.schema.models import FieldDefinition, Record

if TYPE_CHECKING:
    from sqlalchemy import Engine


def generate_uuid() -> str:
    """Generate a new UUID as string."""
    return str(uuid4())


def utc_now() -> datetime:
    """Get current UTC timestamp."""
    return datetime.now(UTC)


class JSONBQuery:
    """JSON-first query builder for KameleonDB.

    Provides CRUD operations where all field data is stored in a single
    JSON column, providing semantic locality and better agent reasoning.

    Works with both PostgreSQL (JSONB) and SQLite (JSON1).

    For complex queries, use get_schema_context() to get schema information
    and generate SQL queries that can be executed directly against the database.
    """

    def __init__(
        self,
        engine: Engine,
        entity_id: str,
        entity_name: str,
        fields: list[FieldDefinition],
    ) -> None:
        """Initialize JSON query builder.

        Args:
            engine: SQLAlchemy engine
            entity_id: Entity definition ID
            entity_name: Entity name (for error messages)
            fields: List of active field definitions
        """
        self._engine = engine
        self._entity_id = entity_id
        self._entity_name = entity_name
        self._fields = {f.name: f for f in fields}
        self._field_names = set(self._fields.keys())
        # Mapping from logical name to storage column name (for renamed fields)
        self._name_to_column = {f.name: f.column_name for f in fields}
        self._column_to_name = {f.column_name: f.name for f in fields}
        # System columns available on all records
        self._system_columns = {"id", "created_at", "updated_at", "created_by"}
        self._all_columns = self._field_names | self._system_columns
        # Detect dialect
        self._is_postgresql = engine.dialect.name == "postgresql"

    def _get_field(self, name: str) -> FieldDefinition:
        """Get field definition by name."""
        if name not in self._fields:
            raise FieldNotFoundError(name, self._entity_name, sorted(self._field_names))
        return self._fields[name]

    def _serialize_value(self, value: Any, field_type: str) -> Any:
        """Serialize a value for JSON storage.

        Both PostgreSQL JSONB and SQLite JSON handle most types natively.
        We only need special handling for datetime (stored as ISO string).
        """
        if value is None:
            return None

        if field_type == "datetime":
            if isinstance(value, datetime):
                return value.isoformat()
            return str(value)
        elif field_type == "json":
            # Already a dict/list, JSON will handle it
            return value
        elif field_type == "bool":
            return bool(value)
        elif field_type == "int":
            return int(value)
        elif field_type == "float":
            return float(value)
        else:
            # string, text, uuid - store as-is
            return value

    def _deserialize_value(self, value: Any, field_type: str) -> Any:
        """Deserialize a value from JSON.

        JSON returns native types, minimal deserialization needed.
        """
        if value is None:
            return None

        if field_type == "datetime":
            # Stored as ISO string, return as-is (or parse if needed)
            return value
        elif field_type == "bool":
            return bool(value)
        elif field_type == "int":
            return int(value) if value is not None else None
        elif field_type == "float":
            return float(value) if value is not None else None
        else:
            return value

    def _record_to_dict(self, record: Record) -> dict[str, Any]:
        """Convert a record to a dictionary."""
        result: dict[str, Any] = {
            "id": record.id,
            "created_at": record.created_at.isoformat() if record.created_at else None,
            "updated_at": record.updated_at.isoformat() if record.updated_at else None,
            "created_by": record.created_by,
        }

        # Merge JSON data, translating column names to logical names
        if record.data:
            for column_name, value in record.data.items():
                # Translate column name to logical name (handles renamed fields)
                logical_name = self._column_to_name.get(column_name, column_name)
                if logical_name in self._field_names:
                    result[logical_name] = value

        # Add None for fields without values
        for field_name in self._field_names:
            if field_name not in result:
                result[field_name] = None

        return result

    def _validate_fields(self, data: dict[str, Any]) -> None:
        """Validate that all fields in data exist."""
        for field in data:
            if field not in self._all_columns:
                raise FieldNotFoundError(field, self._entity_name, sorted(self._field_names))

    def _serialize_record_data(self, data: dict[str, Any]) -> dict[str, Any]:
        """Serialize all field data for JSON storage.

        Uses column_name (storage key) instead of logical name for JSON keys.
        This supports field renaming where the logical name changes but the
        storage key remains the same.
        """
        json_data = {}
        for field_name, value in data.items():
            if field_name in self._system_columns:
                continue  # Skip system columns
            field = self._get_field(field_name)
            # Use column_name for storage (handles renamed fields)
            column_name = self._name_to_column.get(field_name, field_name)
            json_data[column_name] = self._serialize_value(value, field.field_type)
        return json_data

    def insert(
        self,
        data: dict[str, Any],
        created_by: str | None = None,
    ) -> str:
        """Insert a new record.

        Args:
            data: Record data (field: value pairs)
            created_by: Who created this record

        Returns:
            The new record ID
        """
        self._validate_fields(data)

        record_id = generate_uuid()
        now = utc_now()

        try:
            with Session(self._engine) as session:
                # Serialize all field data to JSON
                json_data = self._serialize_record_data(data)

                # Create the record with JSON data
                record = Record(
                    id=record_id,
                    entity_id=self._entity_id,
                    data=json_data,  # All fields in one JSON column
                    created_at=now,
                    updated_at=now,
                    created_by=created_by,
                    is_deleted=False,
                )
                session.add(record)
                session.commit()

            return record_id
        except FieldNotFoundError:
            raise
        except Exception as e:
            raise QueryError(f"Failed to insert record: {e}") from e

    def insert_many(
        self,
        records: list[dict[str, Any]],
        created_by: str | None = None,
    ) -> list[str]:
        """Insert multiple records.

        Args:
            records: List of record data dicts
            created_by: Who created these records

        Returns:
            List of new record IDs
        """
        if not records:
            return []

        # Validate all records first
        for rec in records:
            self._validate_fields(rec)

        record_ids = []
        now = utc_now()

        try:
            with Session(self._engine) as session:
                for record_data in records:
                    record_id = generate_uuid()
                    record_ids.append(record_id)

                    # Serialize all field data to JSON
                    json_data = self._serialize_record_data(record_data)

                    # Create the record
                    record = Record(
                        id=record_id,
                        entity_id=self._entity_id,
                        data=json_data,
                        created_at=now,
                        updated_at=now,
                        created_by=created_by,
                        is_deleted=False,
                    )
                    session.add(record)

                session.commit()

            return record_ids
        except FieldNotFoundError:
            raise
        except Exception as e:
            raise QueryError(f"Failed to insert records: {e}") from e

    def find_by_id(self, record_id: str) -> dict[str, Any] | None:
        """Find a record by ID.

        Args:
            record_id: Record ID

        Returns:
            Record dict or None if not found
        """
        try:
            with Session(self._engine) as session:
                record = (
                    session.query(Record)
                    .filter(Record.id == record_id)
                    .filter(Record.entity_id == self._entity_id)
                    .filter(Record.is_deleted == False)  # noqa: E712
                    .first()
                )

                if not record:
                    return None

                return self._record_to_dict(record)
        except Exception as e:
            raise QueryError(f"Failed to find record: {e}") from e

    def update(
        self,
        record_id: str,
        data: dict[str, Any],
    ) -> dict[str, Any]:
        """Update a record.

        Args:
            record_id: Record ID to update
            data: Fields to update

        Returns:
            Updated record dict
        """
        self._validate_fields(data)

        # Check record exists
        existing = self.find_by_id(record_id)
        if not existing:
            raise RecordNotFoundError(record_id, self._entity_name)

        now = utc_now()

        try:
            with Session(self._engine) as session:
                # Serialize new field data
                json_updates = self._serialize_record_data(data)

                if self._is_postgresql:
                    # Use PostgreSQL || operator to merge JSONB (efficient)
                    session.execute(
                        update(Record)
                        .where(Record.id == record_id)
                        .values(
                            data=Record.data.op("||")(cast(json_updates, JSONB)),
                            updated_at=now,
                        )
                    )
                else:
                    # For SQLite, fetch existing data, merge in Python, and save
                    # This is less efficient but works with SQLite's JSON1
                    record = session.query(Record).filter(Record.id == record_id).first()
                    if record:
                        merged_data = {**(record.data or {}), **json_updates}
                        # Use ORM to update - this handles JSON serialization properly
                        record.data = merged_data
                        record.updated_at = now
                        session.add(record)

                session.commit()

            # Return updated record
            return self.find_by_id(record_id) or {}
        except FieldNotFoundError:
            raise
        except RecordNotFoundError:
            raise
        except Exception as e:
            raise QueryError(f"Failed to update record: {e}") from e

    def delete(self, record_id: str) -> bool:
        """Delete a record (soft delete).

        Args:
            record_id: Record ID to delete

        Returns:
            True if deleted
        """
        # Check record exists
        existing = self.find_by_id(record_id)
        if not existing:
            raise RecordNotFoundError(record_id, self._entity_name)

        try:
            with Session(self._engine) as session:
                # Soft delete the record
                session.execute(
                    update(Record)
                    .where(Record.id == record_id)
                    .values(is_deleted=True, updated_at=utc_now())
                )

                session.commit()

            return True
        except Exception as e:
            raise QueryError(f"Failed to delete record: {e}") from e
