"""SQLite-based session memory for the companion."""

from __future__ import annotations

import json
import sqlite3
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any


@dataclass
class Message:
    """A single message in the conversation."""

    role: str  # "user" | "assistant" | "system"
    content: str
    timestamp: datetime
    metadata: dict[str, Any] | None = None


class SessionMemory:
    """Manages conversation history in SQLite."""

    def __init__(self, db_path: str | Path):
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_db()

    def _init_db(self) -> None:
        """Initialize the database schema."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT NOT NULL,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    metadata TEXT,
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id)
                )
            """)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_messages_session 
                ON messages(session_id, timestamp)
            """)
            conn.commit()

    def get_or_create_session(self, session_id: str) -> str:
        """Get existing session or create new one."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "INSERT OR IGNORE INTO sessions (session_id) VALUES (?)", (session_id,)
            )
            conn.execute(
                """UPDATE sessions SET updated_at = CURRENT_TIMESTAMP 
                   WHERE session_id = ?""",
                (session_id,),
            )
            conn.commit()
        return session_id

    def add_message(
        self,
        session_id: str,
        role: str,
        content: str,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Add a message to the session."""
        self.get_or_create_session(session_id)
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                """INSERT INTO messages (session_id, role, content, metadata)
                   VALUES (?, ?, ?, ?)""",
                (session_id, role, content, json.dumps(metadata) if metadata else None),
            )
            conn.commit()

    def get_messages(
        self, session_id: str, limit: int = 20, before_id: int | None = None
    ) -> list[Message]:
        """Get messages from a session, most recent first."""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            if before_id:
                rows = conn.execute(
                    """SELECT role, content, timestamp, metadata
                       FROM messages 
                       WHERE session_id = ? AND id < ?
                       ORDER BY timestamp DESC
                       LIMIT ?""",
                    (session_id, before_id, limit),
                ).fetchall()
            else:
                rows = conn.execute(
                    """SELECT role, content, timestamp, metadata
                       FROM messages 
                       WHERE session_id = ?
                       ORDER BY timestamp DESC
                       LIMIT ?""",
                    (session_id, limit),
                ).fetchall()

            messages = []
            for row in rows:
                meta = json.loads(row["metadata"]) if row["metadata"] else None
                messages.append(
                    Message(
                        role=row["role"],
                        content=row["content"],
                        timestamp=datetime.fromisoformat(row["timestamp"]),
                        metadata=meta,
                    )
                )
            # Return in chronological order
            return list(reversed(messages))

    def get_session_summary(self, session_id: str) -> dict[str, Any] | None:
        """Get summary info about a session."""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute(
                """SELECT session_id, created_at, updated_at,
                          (SELECT COUNT(*) FROM messages WHERE session_id = ?) as message_count
                   FROM sessions WHERE session_id = ?""",
                (session_id, session_id),
            ).fetchone()
            if not row:
                return None
            return {
                "session_id": row["session_id"],
                "created_at": row["created_at"],
                "updated_at": row["updated_at"],
                "message_count": row["message_count"],
            }

    def list_sessions(self, limit: int = 100) -> list[dict[str, Any]]:
        """List recent sessions."""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(
                """SELECT session_id, created_at, updated_at
                   FROM sessions
                   ORDER BY updated_at DESC
                   LIMIT ?""",
                (limit,),
            ).fetchall()
            return [
                {
                    "session_id": r["session_id"],
                    "created_at": r["created_at"],
                    "updated_at": r["updated_at"],
                }
                for r in rows
            ]

    def clear_session(self, session_id: str) -> None:
        """Clear all messages from a session."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
            conn.commit()

    def delete_session(self, session_id: str) -> None:
        """Delete a session and all its messages."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
            conn.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
            conn.commit()
