"""

Database Management for Swarm Portal



Provides SQLite-based persistence for:

- Chat sessions and messages

- Search results and agent outputs

- Session history and exports



Migration path to Postgres available for production.

"""



import sqlite3

import json

import uuid

from datetime import datetime

from pathlib import Path

from typing import Dict, List, Any, Optional

import logging



logger = logging.getLogger(__name__)





class SwarmPortalDB:

    """

    SQLite database manager for Swarm Portal.



    Handles sessions, messages, and search results with automatic schema creation.

    """



    def __init__(self, db_path: str = "swarm_portal.db"):

        """

        Initialize database connection.



        Args:

            db_path: Path to SQLite database file

        """

        self.db_path = Path(db_path)

        self.conn = None

        self._init_db()



    def _init_db(self):

        """Initialize database connection and create tables if needed"""

        self.conn = sqlite3.connect(str(self.db_path), check_same_thread=False)

        self.conn.row_factory = sqlite3.Row  # Enable dict-like access

        self._create_tables()

        logger.info(f"Database initialized at {self.db_path}")



    def _create_tables(self):

        """Create database tables if they don't exist"""

        cursor = self.conn.cursor()



        # Sessions table

        cursor.execute("""

            CREATE TABLE IF NOT EXISTS sessions (

                id TEXT PRIMARY KEY,

                type TEXT NOT NULL,  -- 'chat', 'search', 'hive'

                user_id TEXT,

                provider TEXT,

                model TEXT,

                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,

                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,

                metadata TEXT  -- JSON

            )

        """)



        # Messages table (for chat sessions)

        cursor.execute("""

            CREATE TABLE IF NOT EXISTS messages (

                id TEXT PRIMARY KEY,

                session_id TEXT NOT NULL,

                role TEXT NOT NULL,  -- 'user', 'assistant', 'tool'

                content TEXT,

                metadata TEXT,  -- JSON

                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,

                FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE

            )

        """)



        # Search results table (for search sessions)

        cursor.execute("""

            CREATE TABLE IF NOT EXISTS search_results (

                id TEXT PRIMARY KEY,

                session_id TEXT NOT NULL,

                agent_id TEXT,

                agent_type TEXT,

                task TEXT,

                content TEXT,

                status TEXT,

                metadata TEXT,  -- JSON

                started_at TIMESTAMP,

                completed_at TIMESTAMP,

                FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE

            )

        """)



        # Create indexes for performance

        cursor.execute("""

            CREATE INDEX IF NOT EXISTS idx_sessions_type

            ON sessions(type)

        """)



        cursor.execute("""

            CREATE INDEX IF NOT EXISTS idx_sessions_created

            ON sessions(created_at)

        """)



        cursor.execute("""

            CREATE INDEX IF NOT EXISTS idx_messages_session

            ON messages(session_id)

        """)



        cursor.execute("""

            CREATE INDEX IF NOT EXISTS idx_search_session

            ON search_results(session_id)

        """)



        self.conn.commit()

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



    def create_session(self, session_type: str, user_id: str = "anonymous",

                      provider: str = None, model: str = None,

                      metadata: Dict = None) -> str:

        """

        Create a new session.



        Args:

            session_type: Type of session ('chat', 'search', 'hive')

            user_id: User identifier

            provider: Provider name

            model: Model name

            metadata: Additional metadata



        Returns:

            Session ID (UUID)

        """

        session_id = str(uuid.uuid4())

        cursor = self.conn.cursor()



        cursor.execute("""

            INSERT INTO sessions (id, type, user_id, provider, model, metadata)

            VALUES (?, ?, ?, ?, ?, ?)

        """, (

            session_id,

            session_type,

            user_id,

            provider,

            model,

            json.dumps(metadata) if metadata else None

        ))



        self.conn.commit()

        logger.info(f"Created {session_type} session: {session_id}")

        return session_id



    def add_message(self, session_id: str, role: str, content: str,

                   metadata: Dict = None) -> str:

        """

        Add a message to a chat session.



        Args:

            session_id: Session ID

            role: Message role ('user', 'assistant', 'tool')

            content: Message content

            metadata: Additional metadata



        Returns:

            Message ID (UUID)

        """

        message_id = str(uuid.uuid4())

        cursor = self.conn.cursor()



        cursor.execute("""

            INSERT INTO messages (id, session_id, role, content, metadata)

            VALUES (?, ?, ?, ?, ?)

        """, (

            message_id,

            session_id,

            role,

            content,

            json.dumps(metadata) if metadata else None

        ))



        # Update session updated_at

        cursor.execute("""

            UPDATE sessions SET updated_at = CURRENT_TIMESTAMP WHERE id = ?

        """, (session_id,))



        self.conn.commit()

        return message_id



    def add_search_result(self, session_id: str, agent_id: str, agent_type: str,

                         task: str, content: str, status: str,

                         metadata: Dict = None) -> str:

        """

        Add a search result to a search session.



        Args:

            session_id: Session ID

            agent_id: Agent identifier

            agent_type: Agent type

            task: Agent task description

            content: Search result content

            status: Status ('pending', 'running', 'completed', 'error')

            metadata: Additional metadata



        Returns:

            Result ID (UUID)

        """

        result_id = str(uuid.uuid4())

        cursor = self.conn.cursor()



        cursor.execute("""

            INSERT INTO search_results

            (id, session_id, agent_id, agent_type, task, content, status, metadata, started_at)

            VALUES (?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)

        """, (

            result_id,

            session_id,

            agent_id,

            agent_type,

            task,

            content,

            status,

            json.dumps(metadata) if metadata else None

        ))



        # Update session updated_at

        cursor.execute("""

            UPDATE sessions SET updated_at = CURRENT_TIMESTAMP WHERE id = ?

        """, (session_id,))



        self.conn.commit()

        return result_id



    def get_session(self, session_id: str) -> Optional[Dict]:

        """Get session by ID"""

        cursor = self.conn.cursor()

        cursor.execute("SELECT * FROM sessions WHERE id = ?", (session_id,))

        row = cursor.fetchone()



        if not row:

            return None



        session = dict(row)

        if session.get('metadata'):

            session['metadata'] = json.loads(session['metadata'])



        return session



    def get_session_messages(self, session_id: str) -> List[Dict]:

        """Get all messages for a session"""

        cursor = self.conn.cursor()

        cursor.execute("""

            SELECT * FROM messages

            WHERE session_id = ?

            ORDER BY timestamp ASC

        """, (session_id,))



        messages = []

        for row in cursor.fetchall():

            msg = dict(row)

            if msg.get('metadata'):

                msg['metadata'] = json.loads(msg['metadata'])

            messages.append(msg)



        return messages



    def get_search_results(self, session_id: str) -> List[Dict]:

        """Get all search results for a session"""

        cursor = self.conn.cursor()

        cursor.execute("""

            SELECT * FROM search_results

            WHERE session_id = ?

            ORDER BY started_at ASC

        """, (session_id,))



        results = []

        for row in cursor.fetchall():

            result = dict(row)

            if result.get('metadata'):

                result['metadata'] = json.loads(result['metadata'])

            results.append(result)



        return results



    def list_sessions(self, session_type: str = None, user_id: str = None,

                     limit: int = 50, offset: int = 0) -> List[Dict]:

        """

        List sessions with optional filtering.



        Args:

            session_type: Filter by type

            user_id: Filter by user

            limit: Max results

            offset: Pagination offset



        Returns:

            List of session dictionaries

        """

        cursor = self.conn.cursor()



        query = "SELECT * FROM sessions WHERE 1=1"

        params = []



        if session_type:

            query += " AND type = ?"

            params.append(session_type)



        if user_id:

            query += " AND user_id = ?"

            params.append(user_id)



        query += " ORDER BY updated_at DESC LIMIT ? OFFSET ?"

        params.extend([limit, offset])



        cursor.execute(query, params)



        sessions = []

        for row in cursor.fetchall():

            session = dict(row)

            if session.get('metadata'):

                session['metadata'] = json.loads(session['metadata'])

            sessions.append(session)



        return sessions



    def delete_session(self, session_id: str):

        """Delete a session and all related data"""

        cursor = self.conn.cursor()

        cursor.execute("DELETE FROM sessions WHERE id = ?", (session_id,))

        self.conn.commit()

        logger.info(f"Deleted session: {session_id}")



    def close(self):

        """Close database connection"""

        if self.conn:

            self.conn.close()

            logger.info("Database connection closed")





# Global database instance

_db = None





def get_db(db_path: str = "swarm_portal.db") -> SwarmPortalDB:

    """

    Get global database instance.



    Args:

        db_path: Path to database file



    Returns:

        SwarmPortalDB instance

    """

    global _db



    if _db is None:

        _db = SwarmPortalDB(db_path)



    return _db





# Convenience exports

__all__ = [

    'SwarmPortalDB',

    'get_db',

]



