# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""The v1 database schema for the DatabaseSessionService.

This module defines SQLAlchemy models for storing session and event data
in a relational database with the "events" table using JSON
serialization for Event data.

See https://github.com/google/adk-python/discussions/3605 for more details.
"""

from __future__ import annotations

from datetime import datetime
from datetime import timezone
from typing import Any

from google.adk.platform import uuid as platform_uuid
from sqlalchemy import desc
from sqlalchemy import ForeignKeyConstraint
from sqlalchemy import func
from sqlalchemy import Index
from sqlalchemy import inspect
from sqlalchemy.ext.mutable import MutableDict
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column
from sqlalchemy.orm import relationship
from sqlalchemy.types import String

from ...events.event import Event
from ..session import Session
from .shared import DEFAULT_MAX_KEY_LENGTH
from .shared import DEFAULT_MAX_VARCHAR_LENGTH
from .shared import DynamicJSON
from .shared import PreciseTimestamp


class Base(DeclarativeBase):
  """Base class for v1 database tables."""

  pass


class StorageMetadata(Base):
  """Represents ADK internal metadata stored in the database.

  This table is used to store internal information like the schema version.
  The DatabaseSessionService will populate and utilize this table to manage
  database compatibility and migrations.
  """

  __tablename__ = "adk_internal_metadata"
  key: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  value: Mapped[str] = mapped_column(String(DEFAULT_MAX_VARCHAR_LENGTH))


class StorageSession(Base):
  """Represents a session stored in the database."""

  __tablename__ = "sessions"

  app_name: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  user_id: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  id: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH),
      primary_key=True,
      default=platform_uuid.new_uuid,
  )

  state: Mapped[MutableDict[str, Any]] = mapped_column(
      MutableDict.as_mutable(DynamicJSON), default={}
  )

  create_time: Mapped[datetime] = mapped_column(
      PreciseTimestamp, default=func.now()
  )
  update_time: Mapped[datetime] = mapped_column(
      PreciseTimestamp, default=func.now(), onupdate=func.now()
  )

  storage_events: Mapped[list[StorageEvent]] = relationship(
      "StorageEvent",
      back_populates="storage_session",
      # Deleting a session will now automatically delete its associated events
      cascade="all, delete-orphan",
  )

  def __repr__(self):
    return f"<StorageSession(id={self.id}, update_time={self.update_time})>"

  @property
  def update_timestamp_tz(self) -> float:
    """Returns the update timestamp as a POSIX timestamp.

    This is a compatibility alias for callers that used the pre-`main` API.
    """
    sqlalchemy_session = inspect(self).session
    is_sqlite = bool(
        sqlalchemy_session
        and sqlalchemy_session.bind
        and sqlalchemy_session.bind.dialect.name == "sqlite"
    )
    is_postgresql = bool(
        sqlalchemy_session
        and sqlalchemy_session.bind
        and sqlalchemy_session.bind.dialect.name == "postgresql"
    )
    return self.get_update_timestamp(
        is_sqlite=is_sqlite, is_postgresql=is_postgresql
    )

  def get_update_timestamp(
      self, is_sqlite: bool = False, is_postgresql: bool = False
  ) -> float:
    """Returns the time zone aware update timestamp."""
    del is_sqlite, is_postgresql  # Unused.
    if self.update_time.tzinfo is None:
      # SQLite and PostgreSQL do not support timezone. SQLAlchemy returns a naive datetime
      # object without timezone information. We need to convert it to UTC
      # manually.
      return self.update_time.replace(tzinfo=timezone.utc).timestamp()
    return self.update_time.timestamp()

  def get_update_marker(self) -> str:
    """Returns a stable revision marker for optimistic concurrency checks."""
    update_time = self.update_time
    if update_time.tzinfo is not None:
      update_time = update_time.astimezone(timezone.utc)
    return update_time.isoformat(timespec="microseconds")

  def to_session(
      self,
      state: dict[str, Any] | None = None,
      events: list[Event] | None = None,
      is_sqlite: bool = False,
      is_postgresql: bool = False,
  ) -> Session:
    """Converts the storage session to a session object."""
    if state is None:
      state = {}
    if events is None:
      events = []

    session = Session(
        app_name=self.app_name,
        user_id=self.user_id,
        id=self.id,
        state=state,
        events=events,
        last_update_time=self.get_update_timestamp(
            is_sqlite=is_sqlite, is_postgresql=is_postgresql
        ),
    )
    session._storage_update_marker = self.get_update_marker()
    return session


class StorageEvent(Base):
  """Represents an event stored in the database."""

  __tablename__ = "events"

  id: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  app_name: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  user_id: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  session_id: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )

  invocation_id: Mapped[str] = mapped_column(String(DEFAULT_MAX_VARCHAR_LENGTH))
  timestamp: Mapped[PreciseTimestamp] = mapped_column(
      PreciseTimestamp, default=func.now()
  )
  # The event_data uses JSON serialization to store the Event data, replacing
  # various fields previously used.
  event_data: Mapped[dict[str, Any]] = mapped_column(DynamicJSON, nullable=True)

  storage_session: Mapped[StorageSession] = relationship(
      "StorageSession",
      back_populates="storage_events",
  )

  __table_args__ = (
      ForeignKeyConstraint(
          ["app_name", "user_id", "session_id"],
          ["sessions.app_name", "sessions.user_id", "sessions.id"],
          ondelete="CASCADE",
      ),
      Index(
          "idx_events_app_user_session_ts",
          "app_name",
          "user_id",
          "session_id",
          desc("timestamp"),
      ),
  )

  @classmethod
  def from_event(cls, session: Session, event: Event) -> StorageEvent:
    """Creates a StorageEvent from an Event."""
    return StorageEvent(
        id=event.id,
        invocation_id=event.invocation_id,
        session_id=session.id,
        app_name=session.app_name,
        user_id=session.user_id,
        timestamp=datetime.fromtimestamp(event.timestamp),
        event_data=event.model_dump(exclude_none=True, mode="json"),
    )

  def to_event(self) -> Event:
    """Converts the StorageEvent to an Event."""
    return Event.model_validate({
        **self.event_data,
        "id": self.id,
        "invocation_id": self.invocation_id,
        "timestamp": self.timestamp.timestamp(),
    })


class StorageAppState(Base):
  """Represents an app state stored in the database."""

  __tablename__ = "app_states"

  app_name: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  state: Mapped[MutableDict[str, Any]] = mapped_column(
      MutableDict.as_mutable(DynamicJSON), default={}
  )
  update_time: Mapped[datetime] = mapped_column(
      PreciseTimestamp, default=func.now(), onupdate=func.now()
  )


class StorageUserState(Base):
  """Represents a user state stored in the database."""

  __tablename__ = "user_states"

  app_name: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  user_id: Mapped[str] = mapped_column(
      String(DEFAULT_MAX_KEY_LENGTH), primary_key=True
  )
  state: Mapped[MutableDict[str, Any]] = mapped_column(
      MutableDict.as_mutable(DynamicJSON), default={}
  )
  update_time: Mapped[datetime] = mapped_column(
      PreciseTimestamp, default=func.now(), onupdate=func.now()
  )
