import asyncpg
import os
from typing import Optional

pool: Optional[asyncpg.Pool] = None

async def init_db():
    global pool
    pool = await asyncpg.create_pool(os.getenv("DATABASE_URL"), min_size=2, max_size=10)
    await _create_tables()

async def get_pool() -> asyncpg.Pool:
    return pool

async def _create_tables():
    async with pool.acquire() as conn:
        await conn.execute("""
        CREATE TABLE IF NOT EXISTS memories (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            user_id TEXT DEFAULT 'default',
            content TEXT NOT NULL,
            embedding vector(768),
            importance FLOAT DEFAULT 0.7,
            access_count INT DEFAULT 0,
            last_accessed TIMESTAMPTZ DEFAULT NOW(),
            created_at TIMESTAMPTZ DEFAULT NOW(),
            memory_type TEXT DEFAULT 'fact',
            metadata JSONB DEFAULT '{}'
        );
        CREATE TABLE IF NOT EXISTS code_nodes (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            project_id TEXT NOT NULL,
            node_type TEXT NOT NULL,
            name TEXT NOT NULL,
            full_name TEXT,
            content TEXT,
            embedding vector(768),
            file_path TEXT,
            line_start INT,
            line_end INT,
            metadata JSONB DEFAULT '{}'
        );
        CREATE TABLE IF NOT EXISTS code_edges (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            project_id TEXT NOT NULL,
            source_id UUID REFERENCES code_nodes(id) ON DELETE CASCADE,
            target_id UUID REFERENCES code_nodes(id) ON DELETE CASCADE,
            relation_type TEXT
        );
        CREATE TABLE IF NOT EXISTS conversations (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            session_id TEXT NOT NULL,
            role TEXT NOT NULL,
            content TEXT NOT NULL,
            embedding vector(768),
            created_at TIMESTAMPTZ DEFAULT NOW()
        );
        CREATE TABLE IF NOT EXISTS conversation_summaries (
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            session_id TEXT NOT NULL,
            summary TEXT NOT NULL,
            turn_start INT DEFAULT 0,
            turn_end INT DEFAULT 0,
            created_at TIMESTAMPTZ DEFAULT NOW()
        );
        CREATE INDEX IF NOT EXISTS idx_mem_embed   ON memories      USING hnsw (embedding vector_cosine_ops);
        CREATE INDEX IF NOT EXISTS idx_node_embed  ON code_nodes    USING hnsw (embedding vector_cosine_ops);
        CREATE INDEX IF NOT EXISTS idx_conv_sess   ON conversations (session_id, created_at DESC);
        CREATE INDEX IF NOT EXISTS idx_mem_user    ON memories      (user_id, last_accessed DESC);
        CREATE INDEX IF NOT EXISTS idx_node_proj   ON code_nodes    (project_id, node_type);
        CREATE INDEX IF NOT EXISTS idx_edge_src    ON code_edges    (source_id);
        CREATE INDEX IF NOT EXISTS idx_edge_dst    ON code_edges    (target_id);
        """)