#!/usr/bin/env python3
"""Initialize Letta database schema by creating all tables."""

import asyncio
import os
from sqlalchemy import create_engine

# Set the database URI
os.environ["LETTA_PG_URI"] = "postgresql://letta:genesis2025@localhost:5432/letta"

# Import the database utilities
from letta.database_utils import get_database_uri_for_context
from letta.settings import settings

# Import all ORM models to ensure they're registered with the Base
from letta.orm.base import Base
from letta.orm import (
    organization,
    user,
    agent,
    message,
    block,
    source,
    tool,
    file,
    job,
    provider,
    provider_model,
    sandbox_config,
    identity,
    group,
    passage,
    passage_tag,
    archive,
    agents_tags,
    archives_agents,
    blocks_agents,
    files_agents,
    groups_agents,
    groups_blocks,
    identities_agents,
    identities_blocks,
    sources_agents,
    tools_agents,
    job_messages,
    llm_batch_items,
    llm_batch_job,
    mcp_oauth,
    mcp_server,
    provider_trace,
    run,
    run_metrics,
    step,
    step_metrics,
    prompt,
    block_history,
)

def main():
    """Create all database tables."""
    print(f"Using database URI: {settings.letta_pg_uri}")

    # Convert to sync URI for schema creation
    sync_uri = get_database_uri_for_context(settings.letta_pg_uri, "sync")
    print(f"Sync URI: {sync_uri}")

    # Create sync engine
    engine = create_engine(sync_uri)

    print("Creating all tables...")
    Base.metadata.create_all(engine)

    print("✅ Database schema initialized successfully!")
    print("\nCreated tables:")
    for table_name in sorted(Base.metadata.tables.keys()):
        print(f"  - {table_name}")

if __name__ == "__main__":
    main()