import asyncpg
import logging
from typing import Optional, List, Dict, Any
from config import config

logger = logging.getLogger(__name__)


class Database:
    def __init__(self) -> None:
        self.pool: Optional[asyncpg.Pool] = None

    async def connect(self) -> None:
        # Render free tier: 512 MB RAM, keep pool small
        self.pool = await asyncpg.create_pool(
            dsn=config.DATABASE_URL,
            min_size=1,
            max_size=3,
            command_timeout=60,
        )
        logger.info("Database pool created (max_size=3)")
        await self._create_tables()

    async def disconnect(self) -> None:
        if self.pool:
            await self.pool.close()
            logger.info("Database pool closed")

    def _acquire(self):
        if self.pool is None:
            raise RuntimeError("Database not connected. Call connect() first.")
        return self.pool.acquire()

    async def _create_tables(self) -> None:
        async with self._acquire() as conn:
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    id BIGINT PRIMARY KEY,
                    username VARCHAR(255),
                    first_name VARCHAR(255),
                    last_name VARCHAR(255),
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                )
            """)
              await conn.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    id SERIAL PRIMARY KEY,
                    user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
                    role VARCHAR(20) NOT NULL CHECK (role IN ('user', 'assistant', 'system')),
                    content TEXT NOT NULL,
                    is_summarized BOOLEAN DEFAULT FALSE,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                )
            """)
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_messages_user_id_created_at
                ON messages(user_id, created_at DESC)
            """)
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_messages_user_id_summarized
                ON messages(user_id, is_summarized, created_at DESC)
            """)
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS summaries (
                    id SERIAL PRIMARY KEY,
                    user_id BIGINT NOT NULL UNIQUE REFERENCES users(id) ON DELETE CASCADE,
                    summary TEXT NOT NULL,
                    message_count INTEGER NOT NULL DEFAULT 0,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                )
            """)

            logger.info("Database tables created/verified")

    async def upsert_user(
        self,
        user_id: int,
        username: Optional[str],
        first_name: Optional[str],
        last_name: Optional[str],
    ) -> None:
        async with self._acquire() as conn:
            await conn.execute("""
                INSERT INTO users (id, username, first_name, last_name)
                VALUES ($1, $2, $3, $4)
                ON CONFLICT (id) DO UPDATE SET
                    username = EXCLUDED.username,
                    first_name = EXCLUDED.first_name,
                    last_name = EXCLUDED.last_name,
                    updated_at = NOW()
            """, user_id, username, first_name, last_name)

    async def save_message(self, user_id: int, role: str, content: str) -> None:
        async with self._acquire() as conn:
            await conn.execute("""
                INSERT INTO messages (user_id, role, content)
                VALUES ($1, $2, $3)
            """, user_id, role, content)

    async def get_messages(self, user_id: int, limit: int = 30) -> List[Dict[str, Any]]:
        async with self._acquire() as conn:
            rows = await conn.fetch("""
                SELECT role, content, created_at
                FROM messages
                WHERE user_id = $1 AND is_summarized = FALSE
                ORDER BY created_at DESC
                LIMIT $2
            """, user_id, limit)
            # Reverse to chronological order for the LLM
            return [
                {"role": r["role"], "content": r["content"], "created_at": r["created_at"]}
                for r in reversed(rows)
            ]

    async def get_summary(self, user_id: int) -> Optional[str]:
        async with self._acquire() as conn:
            row = await conn.fetchrow("""
                SELECT summary FROM summaries WHERE user_id = $1
            """, user_id)
            return row["summary"] if row else None

    async def save_summary(self, user_id: int, summary: str, message_count: int) -> None:
        async with self._acquire() as conn:
            await conn.execute("""
                INSERT INTO summaries (user_id, summary, message_count, updated_at)
                VALUES ($1, $2, $3, NOW())
                ON CONFLICT (user_id) DO UPDATE SET
                    summary = EXCLUDED.summary,
                    message_count = summaries.message_count + EXCLUDED.message_count,
                    updated_at = NOW()
            """, user_id, summary, message_count)

    async def mark_summarized(self, user_id: int, cutoff_id: int) -> None:
        async with self._acquire() as conn:
            await conn.execute("""
                UPDATE messages
                SET is_summarized = TRUE
                WHERE user_id = $1 AND id <= $2
            """, user_id, cutoff_id)

    async def get_oldest_unsummarized(self, user_id: int, limit: int) -> List[Dict[str, Any]]:
        async with self._acquire() as conn:
            rows = await conn.fetch("""
                SELECT id, role, content
                FROM messages
                WHERE user_id = $1 AND is_summarized = FALSE
                ORDER BY created_at ASC
                LIMIT $2
            """, user_id, limit)
            return [{"id": r["id"], "role": r["role"], "content": r["content"]} for r in rows]

    async def count_unsummarized(self, user_id: int) -> int:
        async with self._acquire() as conn:
            return await conn.fetchval("""
                SELECT COUNT(*) FROM messages
                WHERE user_id = $1 AND is_summarized = FALSE
            """, user_id) or 0

    async def clear_history(self, user_id: int) -> int:
        async with self._acquire() as conn:
            result = await conn.execute("""
                DELETE FROM messages WHERE user_id = $1
            """, user_id)
            await conn.execute("""
                DELETE FROM summaries WHERE user_id = $1
            """, user_id)
            try:
                count = int(result.split()[-1])
            except (ValueError, IndexError):
                count = 0
            logger.info("Cleared %d messages and summary for user %s", count, user_id)
            return count


    async def get_stats(self, user_id: int) -> Dict[str, Any]:
        async with self._acquire() as conn:
            user_count = await conn.fetchval("SELECT COUNT(*) FROM users")
            msg_count = await conn.fetchval(
                "SELECT COUNT(*) FROM messages WHERE user_id = $1", user_id
            )
            total_msg_count = await conn.fetchval("SELECT COUNT(*) FROM messages")
            return {
                "total_users": user_count,
                "user_messages": msg_count,
                "total_messages": total_msg_count,
            }


db = Database()