"""
Context Memory Store — SQLite storage for Session Nodes and Timeline Events.

PRD §5 — Data Models:
    SessionNode:    session summary + status
    TimelineEvent:  decision / action / result / error events

PRD §6 — Namespace: openmemo/context/{session_id}
"""

import json
import time
import uuid
import sqlite3
import threading
import logging
from dataclasses import dataclass, field, asdict
from typing import List, Optional, Dict

logger = logging.getLogger("openmemo.context_memory")


# ── Data Models ───────────────────────────────────────────────────────────────

@dataclass
class SessionNode:
    """PRD §5.1 — Session Node."""
    session_id: str
    user_id: str = ""
    task_id: str = ""
    summary: str = ""           # compressed summary of the session
    goal: str = ""
    progress: str = ""
    decisions: str = ""
    issues: str = ""
    message_count: int = 0      # total messages in session
    last_updated: float = field(default_factory=time.time)
    created_at: float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_row(cls, row) -> "SessionNode":
        return cls(
            session_id=row["session_id"],
            user_id=row["user_id"],
            task_id=row["task_id"],
            summary=row["summary"],
            goal=row["goal"],
            progress=row["progress"],
            decisions=row["decisions"],
            issues=row["issues"],
            message_count=row["message_count"],
            last_updated=row["last_updated"],
            created_at=row["created_at"],
        )


@dataclass
class TimelineEvent:
    """PRD §5.2 — Timeline Event."""
    session_id: str
    type: str               # "decision" | "action" | "result" | "error"
    content: str
    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_row(cls, row) -> "TimelineEvent":
        return cls(
            event_id=row["event_id"],
            session_id=row["session_id"],
            type=row["type"],
            content=row["content"],
            timestamp=row["timestamp"],
        )


@dataclass
class RawMessage:
    """PRD §5.3 — Raw Context (optional, stored for progressive recall)."""
    session_id: str
    role: str           # "user" | "assistant" | "tool"
    content: str
    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_row(cls, row) -> "RawMessage":
        return cls(
            message_id=row["message_id"],
            session_id=row["session_id"],
            role=row["role"],
            content=row["content"],
            timestamp=row["timestamp"],
        )


# ── Store ─────────────────────────────────────────────────────────────────────

class ContextStore:
    """
    SQLite-backed store for Context Memory.

    Usage:
        store = ContextStore("openmemo.db")
        store.save_session(SessionNode(session_id="s1", user_id="alice",
                                        summary="Deploying FastAPI backend"))
        store.add_event(TimelineEvent(session_id="s1", type="action",
                                       content="Docker image built"))
    """

    def __init__(self, db_path: str = "openmemo.db"):
        self.db_path = db_path
        self._local = threading.local()
        self._init_tables()

    @property
    def conn(self) -> sqlite3.Connection:
        if not hasattr(self._local, "conn") or self._local.conn is None:
            c = sqlite3.connect(self.db_path, check_same_thread=False)
            c.row_factory = sqlite3.Row
            c.execute("PRAGMA journal_mode=WAL")
            c.execute("PRAGMA synchronous=NORMAL")
            self._local.conn = c
        return self._local.conn

    def _init_tables(self):
        self.conn.executescript("""
            CREATE TABLE IF NOT EXISTS context_sessions (
                session_id    TEXT PRIMARY KEY,
                user_id       TEXT DEFAULT '',
                task_id       TEXT DEFAULT '',
                summary       TEXT DEFAULT '',
                goal          TEXT DEFAULT '',
                progress      TEXT DEFAULT '',
                decisions     TEXT DEFAULT '',
                issues        TEXT DEFAULT '',
                message_count INTEGER DEFAULT 0,
                last_updated  REAL NOT NULL,
                created_at    REAL NOT NULL
            );

            CREATE TABLE IF NOT EXISTS context_timeline (
                event_id    TEXT PRIMARY KEY,
                session_id  TEXT NOT NULL,
                type        TEXT NOT NULL,
                content     TEXT NOT NULL,
                timestamp   REAL NOT NULL
            );

            CREATE INDEX IF NOT EXISTS idx_ctx_session
                ON context_sessions(user_id);
            CREATE INDEX IF NOT EXISTS idx_ctx_timeline
                ON context_timeline(session_id, timestamp);

            CREATE TABLE IF NOT EXISTS context_raw_messages (
                message_id  TEXT PRIMARY KEY,
                session_id  TEXT NOT NULL,
                role        TEXT NOT NULL,
                content     TEXT NOT NULL,
                timestamp   REAL NOT NULL
            );

            CREATE INDEX IF NOT EXISTS idx_ctx_raw
                ON context_raw_messages(session_id, timestamp);
        """)
        self.conn.commit()

    # ── Session ────────────────────────────────────────────────────────────────

    def save_session(self, session: SessionNode) -> str:
        self.conn.execute("""
            INSERT INTO context_sessions
                (session_id, user_id, task_id, summary, goal, progress,
                 decisions, issues, message_count, last_updated, created_at)
            VALUES (?,?,?,?,?,?,?,?,?,?,?)
            ON CONFLICT(session_id) DO UPDATE SET
                summary=excluded.summary, goal=excluded.goal,
                progress=excluded.progress, decisions=excluded.decisions,
                issues=excluded.issues, message_count=excluded.message_count,
                last_updated=excluded.last_updated
        """, (
            session.session_id, session.user_id, session.task_id,
            session.summary, session.goal, session.progress,
            session.decisions, session.issues, session.message_count,
            session.last_updated, session.created_at,
        ))
        self.conn.commit()
        return session.session_id

    def get_session(self, session_id: str) -> Optional[SessionNode]:
        row = self.conn.execute(
            "SELECT * FROM context_sessions WHERE session_id=?", (session_id,)
        ).fetchone()
        return SessionNode.from_row(row) if row else None

    def get_or_create_session(self, session_id: str,
                               user_id: str = "", task_id: str = "") -> SessionNode:
        session = self.get_session(session_id)
        if not session:
            session = SessionNode(session_id=session_id,
                                  user_id=user_id, task_id=task_id)
            self.save_session(session)
        return session

    def list_sessions(self, user_id: str = "", limit: int = 50) -> List[SessionNode]:
        if user_id:
            rows = self.conn.execute(
                "SELECT * FROM context_sessions WHERE user_id=? ORDER BY last_updated DESC LIMIT ?",
                (user_id, limit)
            ).fetchall()
        else:
            rows = self.conn.execute(
                "SELECT * FROM context_sessions ORDER BY last_updated DESC LIMIT ?",
                (limit,)
            ).fetchall()
        return [SessionNode.from_row(r) for r in rows]

    def increment_message_count(self, session_id: str) -> int:
        self.conn.execute(
            "UPDATE context_sessions SET message_count = message_count + 1, last_updated=? WHERE session_id=?",
            (time.time(), session_id)
        )
        self.conn.commit()
        row = self.conn.execute(
            "SELECT message_count FROM context_sessions WHERE session_id=?",
            (session_id,)
        ).fetchone()
        return row["message_count"] if row else 0

    def delete_session(self, session_id: str) -> bool:
        self.conn.execute("DELETE FROM context_sessions WHERE session_id=?", (session_id,))
        self.conn.execute("DELETE FROM context_timeline WHERE session_id=?", (session_id,))
        self.conn.execute("DELETE FROM context_raw_messages WHERE session_id=?", (session_id,))
        self.conn.commit()
        return True

    # ── Timeline ───────────────────────────────────────────────────────────────

    def add_event(self, event: TimelineEvent) -> str:
        self.conn.execute(
            "INSERT INTO context_timeline (event_id, session_id, type, content, timestamp) VALUES (?,?,?,?,?)",
            (event.event_id, event.session_id, event.type, event.content, event.timestamp)
        )
        self.conn.commit()
        return event.event_id

    def get_timeline(self, session_id: str,
                     limit: int = 200) -> List[TimelineEvent]:
        rows = self.conn.execute(
            "SELECT * FROM context_timeline WHERE session_id=? ORDER BY timestamp ASC LIMIT ?",
            (session_id, limit)
        ).fetchall()
        return [TimelineEvent.from_row(r) for r in rows]

    def get_recent_events(self, session_id: str,
                          limit: int = 5) -> List[TimelineEvent]:
        rows = self.conn.execute(
            "SELECT * FROM context_timeline WHERE session_id=? ORDER BY timestamp DESC LIMIT ?",
            (session_id, limit)
        ).fetchall()
        return list(reversed([TimelineEvent.from_row(r) for r in rows]))

    def count_events(self, session_id: str) -> int:
        row = self.conn.execute(
            "SELECT COUNT(*) as n FROM context_timeline WHERE session_id=?",
            (session_id,)
        ).fetchone()
        return row["n"] if row else 0

    # ── Raw Messages ───────────────────────────────────────────────────────────

    def add_raw_message(self, msg: RawMessage) -> str:
        self.conn.execute(
            "INSERT INTO context_raw_messages (message_id, session_id, role, content, timestamp) VALUES (?,?,?,?,?)",
            (msg.message_id, msg.session_id, msg.role, msg.content, msg.timestamp)
        )
        self.conn.commit()
        return msg.message_id

    def get_raw_messages(self, session_id: str,
                         limit: int = 50) -> List[RawMessage]:
        rows = self.conn.execute(
            "SELECT * FROM context_raw_messages WHERE session_id=? ORDER BY timestamp ASC LIMIT ?",
            (session_id, limit)
        ).fetchall()
        return [RawMessage.from_row(r) for r in rows]

    def count_raw_messages(self, session_id: str) -> int:
        row = self.conn.execute(
            "SELECT COUNT(*) as n FROM context_raw_messages WHERE session_id=?",
            (session_id,)
        ).fetchone()
        return row["n"] if row else 0

    def close(self):
        if hasattr(self._local, "conn") and self._local.conn:
            self._local.conn.close()
            self._local.conn = None
