import json
import sqlite3
from typing import Optional

from .base import Session, SessionStore

_SCHEMA = """
CREATE TABLE IF NOT EXISTS sessions (
    user_id TEXT NOT NULL,
    session_id TEXT NOT NULL,
    messages TEXT NOT NULL,    -- JSON-encoded message list
    updated_at REAL NOT NULL,
    PRIMARY KEY (user_id, session_id)
);
"""


class SQLiteSessionStore(SessionStore):
    """A simple SQLite-based session store."""

    def __init__(self, db_path: str):
        self.db_path = db_path
        self._init_schema()
    
    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        conn.row_factory = sqlite3.Row
        return conn

    def _init_schema(self):
        with self._connect() as conn:
            conn.execute(_SCHEMA)

    def load(self, user_id: str, session_id: str) -> Session:
        with self._connect() as conn:
            row: Optional[sqlite3.Row] = conn.execute(
                "SELECT messages FROM agent_session WHERE user_id = ? AND session_id = ?",
                (user_id, session_id),
            ).fetchone()
            messages = []
            if row is not None:
                try:
                    messages = json.loads(row["messages"])
                except (json.JSONDecodeError, TypeError):
                    messages = []
            return Session(session_id=session_id, user_id=user_id, messages=messages)

    def save(self, session: Session) -> None:
        import time
        
        with self._connect() as conn:
            conn.execute(
                """
                INSERT INTO agent_session (user_id, session_id, messages, updated_at)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(user_id, session_id)
                DO UPDATE SET messages=excluded.messages, updated_at=excluded.updated_at
                """,
                (
                    session.user_id,
                    session.session_id,
                    json.dumps(session.messages, ensure_ascii=False),
                    time.time()
                ),
            )
            conn.commit()

    def clear(self, user_id: str, session_id: str) -> None:
        with self._connect() as conn:
            conn.execute(
                "DELETE FROM agent_session WHERE user_id = ? AND session_id = ?",
                (user_id, session_id),
            )
            conn.commit()