db/models.py

import json
import secrets
import sqlite3
import logging
from datetime import datetime, timezone
from pathlib import Path

from db.init import get_conn

log = logging.getLogger("db.models")


def now_iso():
    return datetime.now(timezone.utc).isoformat()


# ============================================================
# USERS
# ============================================================
def get_user(user_id: int):
    with get_conn() as conn:
        row = conn.execute(
            "SELECT * FROM users WHERE user_id = ?", (user_id,)
        ).fetchone()
    return dict(row) if row else None


def upsert_user(user_id: int, username: str = None, first_name: str = None):
    with get_conn() as conn:
        existing = conn.execute(
            "SELECT user_id FROM users WHERE user_id = ?", (user_id,)
        ).fetchone()
        if existing:
            conn.execute(
                "UPDATE users SET username=?, first_name=?, last_seen=? WHERE user_id=?",
                (username, first_name, now_iso(), user_id),
            )
        else:
            conn.execute(
                """INSERT INTO users(user_id, username, first_name, created_at, last_seen)
                   VALUES (?, ?, ?, ?, ?)""",
                (user_id, username, first_name, now_iso(), now_iso()),
            )


def bind_token(user_id: int, token: str) -> bool:
    with get_conn() as conn:
        row = conn.execute(
            "SELECT * FROM tokens WHERE token = ?", (token,)
        ).fetchone()
        if not row:
            return False
        if not row["is_active"]:
            return False
        if row["assigned_to"] not in (None, user_id):
            return False

        conn.execute(
            "UPDATE tokens SET assigned_to=? WHERE token=?", (user_id, token)
        )
        conn.execute(
            "UPDATE users SET token=?, is_active=1 WHERE user_id=?",
            (token, user_id),
        )
        return True


def user_has_valid_token(user_id: int) -> bool:
    with get_conn() as conn:
        row = conn.execute(
            """SELECT u.is_active AS u_active, t.is_active AS t_active
               FROM users u LEFT JOIN tokens t ON u.token = t.token
               WHERE u.user_id = ?""",
            (user_id,),
        ).fetchone()
    if not row:
        return False
    return bool(row["u_active"]) and bool(row["t_active"])


def revoke_user(user_id: int):
    with get_conn() as conn:
        conn.execute("UPDATE users SET is_active=0 WHERE user_id=?", (user_id,))
        conn.execute(
            "UPDATE tokens SET is_active=0 WHERE assigned_to=?", (user_id,)
        )


def delete_user(user_id: int):
    with get_conn() as conn:
        conn.execute("DELETE FROM conversations WHERE user_id=?", (user_id,))
        conn.execute("DELETE FROM users WHERE user_id=?", (user_id,))
        conn.execute(
            "UPDATE tokens SET assigned_to=NULL, is_active=0 WHERE assigned_to=?",
            (user_id,),
        )


def increment_user_turns(user_id: int):
    with get_conn() as conn:
        conn.execute(
            "UPDATE users SET total_turns = total_turns + 1 WHERE user_id=?",
            (user_id,),
        )


# ============================================================
# TOKENS
# ============================================================
def generate_token(assigned_to: int = None) -> str:
    tok = secrets.token_urlsafe(24)
    with get_conn() as conn:
        conn.execute(
            """INSERT INTO tokens(token, assigned_to, created_at)
               VALUES (?, ?, ?)""",
            (tok, assigned_to, now_iso()),
        )
        if assigned_to is not None:
            conn.execute(
                "UPDATE users SET token=?, is_active=1 WHERE user_id=?",
                (tok, assigned_to),
            )
    return tok


def list_tokens():
    with get_conn() as conn:
        rows = conn.execute(
            "SELECT * FROM tokens ORDER BY created_at DESC"
        ).fetchall()
    return [dict(r) for r in rows]


def revoke_token(token: str):
    with get_conn() as conn:
        conn.execute("UPDATE tokens SET is_active=0 WHERE token=?", (token,))
        conn.execute(
            "UPDATE users SET is_active=0 WHERE token=?", (token,)
        )


def reassign_token(token: str, new_user_id: int):
    with get_conn() as conn:
        conn.execute(
            "UPDATE tokens SET assigned_to=? WHERE token=?", (new_user_id, token)
        )
        conn.execute(
            "UPDATE users SET token=?, is_active=1 WHERE user_id=?",
            (token, new_user_id),
        )


def bump_token_use(token: str):
    with get_conn() as conn:
        conn.execute(
            """UPDATE tokens SET turns_used = turns_used + 1,
                                  last_used = ?
               WHERE token = ?""",
            (now_iso(), token),
        )


# ============================================================
# CONVERSATIONS (live context only — cleared by /reset)
# ============================================================
def add_message(user_id: int, role: str, content: str,
                latency_ms: int = None, gpu_time_ms: int = None):
    with get_conn() as conn:
        conn.execute(
            """INSERT INTO conversations
               (user_id, role, content, timestamp, latency_ms, gpu_time_ms)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (user_id, role, content, now_iso(), latency_ms, gpu_time_ms),
        )


def get_history(user_id: int, limit_turns: int = 6):
    """Return last N turns (2*N messages) in chronological order."""
    with get_conn() as conn:
        rows = conn.execute(
            """SELECT role, content FROM conversations
               WHERE user_id = ?
               ORDER BY id DESC LIMIT ?""",
            (user_id, limit_turns * 2),
        ).fetchall()
    rows = list(reversed(rows))
    return [{"role": r["role"], "content": r["content"]} for r in rows]


def clear_history(user_id: int):
    with get_conn() as conn:
        conn.execute("DELETE FROM conversations WHERE user_id=?", (user_id,))


def list_users():
    with get_conn() as conn:
        rows = conn.execute(
            "SELECT * FROM users ORDER BY created_at DESC"
        ).fetchall()
    return [dict(r) for r in rows]


def list_conversations(user_id: int = None, limit: int = 500):
    with get_conn() as conn:
        if user_id:
            rows = conn.execute(
                """SELECT * FROM conversations WHERE user_id=?
                   ORDER BY id DESC LIMIT ?""",
                (user_id, limit),
            ).fetchall()
        else:
            rows = conn.execute(
                "SELECT * FROM conversations ORDER BY id DESC LIMIT ?",
                (limit,),
            ).fetchall()
    return [dict(r) for r in rows]


def stats():
    with get_conn() as conn:
        total_users = conn.execute("SELECT COUNT(*) c FROM users").fetchone()["c"]
        active_users = conn.execute(
            "SELECT COUNT(*) c FROM users WHERE is_active=1"
        ).fetchone()["c"]
        total_msgs = conn.execute(
            "SELECT COUNT(*) c FROM conversations"
        ).fetchone()["c"]
        today = datetime.now(timezone.utc).date().isoformat()
        today_msgs = conn.execute(
            "SELECT COUNT(*) c FROM conversations WHERE timestamp LIKE ?",
            (today + "%",),
        ).fetchone()["c"]
    return {
        "total_users": total_users,
        "active_users": active_users,
        "total_messages": total_msgs,
        "messages_today": today_msgs,
    }


# ============================================================
# SETTINGS
# ============================================================
def get_setting(key: str, default=None):
    with get_conn() as conn:
        row = conn.execute(
            "SELECT value FROM settings WHERE key=?", (key,)
        ).fetchone()
    return row["value"] if row else default


def set_setting(key: str, value: str):
    with get_conn() as conn:
        conn.execute(
            "INSERT INTO settings(key, value) VALUES (?, ?) "
            "ON CONFLICT(key) DO UPDATE SET value=excluded.value",
            (key, str(value)),
        )


def all_settings():
    with get_conn() as conn:
        rows = conn.execute("SELECT key, value FROM settings").fetchall()
    return {r["key"]: r["value"] for r in rows}


# ============================================================
# GPU LOG
# ============================================================
def log_gpu(user_id: int, duration_ms: int, status: str):
    with get_conn() as conn:
        conn.execute(
            """INSERT INTO gpu_log(user_id, timestamp, duration_ms, status)
               VALUES (?, ?, ?, ?)""",
            (user_id, now_iso(), duration_ms, status),
        )


def recent_gpu_log(limit: int = 50):
    with get_conn() as conn:
        rows = conn.execute(
            "SELECT * FROM gpu_log ORDER BY id DESC LIMIT ?", (limit,)
        ).fetchall()
    return [dict(r) for r in rows]


# ============================================================
# ARCHIVE — one markdown file PER SESSION
# Path: /data/chat_logs/{user_id}/{timestamp}.md
# A new session starts whenever the previous one was closed by /reset
# (or when the user hasn't chatted before).
# ============================================================
import os
import re
from datetime import datetime as _dt

ARCHIVE_DIR = os.environ.get("ARCHIVE_DIR", "/data/chat_logs")

# In-process tracker: user_id -> Path of the currently-open session file.
# Reset when /reset is called, so the next message opens a fresh file.
_active_session = {}


def _user_dir(user_id: int) -> Path:
    d = Path(ARCHIVE_DIR) / str(user_id)
    d.mkdir(parents=True, exist_ok=True)
    return d


def _new_session_path(user_id: int) -> Path:
    """Generate a fresh session file path with a UTC timestamp."""
    d = _user_dir(user_id)
    stamp = _dt.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
    p = d / f"{stamp}.md"
    # Extremely unlikely, but avoid clobbering if two messages land in the
    # same second.
    n = 1
    while p.exists():
        p = d / f"{stamp}_{n}.md"
        n += 1
    return p


def archive_append(user_id: int, role: str, content: str,
                   latency_ms: int = None):
    """
    Append a message to the user's *current* session markdown file.
    If no session is open, start a new one.
    """
    p = _active_session.get(user_id)
    if p is None or not p.exists():
        p = _new_session_path(user_id)
        _active_session[user_id] = p
        # Write the file header
        p.write_text(
            f"# Transcript — user {user_id}\n\n"
            f"_Session started: {now_iso()}_\n\n---\n\n",
            encoding="utf-8",
        )

    header = "🧑 **User**" if role == "user" else "🤖 **Assistant**"
    with p.open("a", encoding="utf-8") as f:
        f.write(f"### {header} — `{now_iso()}`\n\n")
        if role == "assistant" and latency_ms is not None:
            f.write(f"_(latency: {latency_ms} ms)_\n\n")
        f.write(content.strip() + "\n\n")


def archive_close_session(user_id: int):
    """
    Close the current session. The next message will open a new file.
    Called by /reset.
    """
    p = _active_session.pop(user_id, None)
    if p and p.exists():
        with p.open("a", encoding="utf-8") as f:
            f.write(f"\n---\n\n_Session closed: {now_iso()}_\n")


def archive_list(user_id: int = None):
    """
    List archived session files.
    If user_id is None, list all users' files.
    Returns list of dicts: {user_id, filename, path, size, modified}
    """
    root = Path(ARCHIVE_DIR)
    if not root.exists():
        return []

    out = []
    user_dirs = [root / str(user_id)] if user_id else [d for d in root.iterdir() if d.is_dir()]
    for d in user_dirs:
        if not d.exists():
            continue
        try:
            uid = int(d.name)
        except ValueError:
            continue
        for f in sorted(d.glob("*.md"), reverse=True):
            stat = f.stat()
            out.append({
                "user_id": uid,
                "filename": f.name,
                "path": str(f),
                "size": stat.st_size,
                "modified": _dt.fromtimestamp(stat.st_mtime, timezone.utc).isoformat(),
            })
    return out


def archive_read(user_id: int, filename: str) -> str:
    p = Path(ARCHIVE_DIR) / str(user_id) / filename
    if not p.exists():
        return ""
    return p.read_text(encoding="utf-8")


def archive_delete(user_id: int, filename: str) -> bool:
    p = Path(ARCHIVE_DIR) / str(user_id) / filename
    if p.exists():
        p.unlink()
        # Close in-process session if we just deleted the active file
        if _active_session.get(user_id) == p:
            _active_session.pop(user_id, None)
        return True
    return False


def archive_delete_user(user_id: int) -> int:
    """Delete ALL session files for one user. Returns count deleted."""
    d = Path(ARCHIVE_DIR) / str(user_id)
    if not d.exists():
        return 0
    n = 0
    for f in d.glob("*.md"):
        f.unlink()
        n += 1
    _active_session.pop(user_id, None)
    return n


def archive_delete_all():
    """Delete every transcript for every user. Returns count deleted."""
    root = Path(ARCHIVE_DIR)
    if not root.exists():
        return 0
    n = 0
    for d in root.iterdir():
        if not d.is_dir():
            continue
        for f in d.glob("*.md"):
            f.unlink()
            n += 1
    _active_session.clear()
    return n


def archive_export_user(user_id: int) -> str:
    """Concatenate all of a user's sessions into one big markdown blob."""
    sessions = archive_list(user_id)
    # Oldest first
    sessions = sorted(sessions, key=lambda x: x["filename"])
    parts = [f"# Full transcript — user {user_id}\n"]
    for s in sessions:
        parts.append(f"\n\n## Session: {s['filename']}\n")
        parts.append(archive_read(user_id, s["filename"]))
    return "".join(parts)