from typing import Optional
from sqlalchemy import create_engine, Column, Integer, JSON, ForeignKey
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy.sql.expression import select

class Base(DeclarativeBase):
    pass

def create_sql_engine (data_directory):
    sql_engine = create_engine(f"sqlite:///{data_directory}/horde.db")
    Base.metadata.create_all(sql_engine)
    return sql_engine

class User(Base):
    __modelname__ = "user"
    __tablename__ = 'users'

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)

    email: Mapped[str]
    username: Mapped[str]
    # Stores the plain password. This must be changed eventually to store a salted hash instead.
    password: Mapped[str]

class Agent(Base):
    __modelname__ = "agent"
    __tablename__ = 'agents'

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    user_id: Mapped[int] = mapped_column(ForeignKey("users.id"))

    agentname: Mapped[Optional[str]]
    visibility: Mapped[int] = mapped_column(default=0)
    # Comma separated list of permission symbols
    permissions: Mapped[str] = mapped_column(default="")
    main_chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id"))

    name: Mapped[str]
    description: Mapped[Optional[str]]
    picture: Mapped[Optional[str]]
    ordering: Mapped[Optional[int]]

    # Agents can be any of a set of agent classes
    kind: Mapped[str]
    # Stores the agent data specific to its class
    data = mapped_column(JSON, default={})

class Chat(Base):
    __modelname__ = 'chat'
    __tablename__ = 'chats'

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    user_id: Mapped[int] = mapped_column(ForeignKey("users.id"))

    name: Mapped[str]
    description: Mapped[Optional[str]]
    picture: Mapped[Optional[str]]
    ordering: Mapped[Optional[int]]

    # For indexing purposes, some information is taken from the last chat message.
    # Preview is the (possibly cropped) content, time is the timestamp of the message.
    last_message_preview: Mapped[Optional[str]]
    last_message_time: Mapped[Optional[int]]

# Stores membership of agents in chats. If a record exists linking an agent and a chat,
# it means the agent participates in that chat.
class Membership(Base):
    __modelname__ = 'membership'
    __tablename__ = 'memberships'

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    agent_id: Mapped[int] = mapped_column(ForeignKey("agents.id"))
    chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id"))

class Message(Base):
    __modelname__ = 'message'
    __tablename__ = 'messages'

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id"))
    thread_id: Mapped[Optional[int]] = mapped_column(ForeignKey("threads.id"))

    # Identifier of the sender: user / system / user-$ID / agent-$ID
    sender: Mapped[str]
    # Timestamp when this message was sent
    time: Mapped[int]
    # Timestamp of this message's last edit
    #last_edit: Mapped[int]

    # Predefined set of message types:
    # message, log, memory, file, note
    kind: Mapped[str]
    # Each kind has a set of subkinds, and a default for empty values
    # message: (empty)
    # log: debug, info
    # memory: (depends on the agent)
    # file: (empty)
    # note: public, private, pinned
    subkind: Mapped[Optional[str]]
    # For the kinds message, log and note, a format can be specified,
    # empty defaults to 'text' and can also be 'markdown', 'csv'
    format: Mapped[Optional[str]]
    content: Mapped[str]

# Represents a temporary topic or stage in a conversation
class Thread(Base):
    __modelname__ = 'thread'
    __tablename__ = 'threads'

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id"))

    # Optional name and descriptions of the thread
    name: Mapped[Optional[str]]
    description: Mapped[Optional[str]]

# A persistent object or value associated with an agent
class AgentResource(Base):
    __modelname__ = 'agent_resource'
    __tablename__ = 'agent_resources'

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    agent_id: Mapped[int] = mapped_column(ForeignKey("agents.id"))

    # Kind is the class of the resource
    # Key is the specific identifier within the class
    # Multiple resources can share key
    kind: Mapped[str]
    key: Mapped[str]
    value = mapped_column(JSON, default=None)