db.py

import sqlite3
from datetime import datetime

DB_PATH = "chatai.db"

def get_conn():
    return sqlite3.connect(DB_PATH, check_same_thread=False)

def init_db():
    conn = get_conn()
    c = conn.cursor()
    c.execute("""
    CREATE TABLE IF NOT EXISTS users (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        username TEXT UNIQUE NOT NULL,
        password_hash BLOB NOT NULL,
        role TEXT NOT NULL DEFAULT 'user',
        is_active INTEGER NOT NULL DEFAULT 1,
        created_at TEXT NOT NULL
    )
    """)
    c.execute("""
    CREATE TABLE IF NOT EXISTS conversations (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        user_id INTEGER NOT NULL,
        title TEXT,
        created_at TEXT NOT NULL,
        FOREIGN KEY(user_id) REFERENCES users(id)
    )
    """)
    c.execute("""
    CREATE TABLE IF NOT EXISTS messages (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        conversation_id INTEGER NOT NULL,
        role TEXT NOT NULL,
        content TEXT NOT NULL,
        attachments TEXT,
        created_at TEXT NOT NULL,
        FOREIGN KEY(conversation_id) REFERENCES conversations(id)
    )
    """)
    conn.commit()
    return conn

# --- User CRUD ---
def create_user(username, password_hash, role="user", is_active=True):
    conn = get_conn()
    c = conn.cursor()
    c.execute(
        "INSERT INTO users (username, password_hash, role, is_active, created_at) VALUES (?,?,?,?,?)",
        (username, password_hash, role, 1 if is_active else 0, datetime.utcnow().isoformat()),
    )
    conn.commit()
    return c.lastrowid

def get_user_by_username(username):
    conn = get_conn()
    c = conn.cursor()
    c.execute("SELECT id, username, password_hash, role, is_active, created_at FROM users WHERE username = ?", (username,))
    row = c.fetchone()
    if not row: return None
    return {"id": row[0], "username": row[1], "password_hash": row[2], "role": row[3], "is_active": bool(row[4]), "created_at": row[5]}

def get_user_by_id(user_id):
    conn = get_conn()
    c = conn.cursor()
    c.execute("SELECT id, username, password_hash, role, is_active, created_at FROM users WHERE id = ?", (user_id,))
    row = c.fetchone()
    if not row: return None
    return {"id": row[0], "username": row[1], "password_hash": row[2], "role": row[3], "is_active": bool(row[4]), "created_at": row[5]}

def list_users():
    conn = get_conn()
    c = conn.cursor()
    c.execute("SELECT id, username, role, is_active, created_at FROM users ORDER BY id ASC")
    rows = c.fetchall()
    return [{"id": r[0], "username": r[1], "role": r[2], "is_active": bool(r[3]), "created_at": r[4]} for r in rows]

def set_user_active(user_id, active: bool):
    conn = get_conn()
    c = conn.cursor()
    c.execute("UPDATE users SET is_active=? WHERE id=?", (1 if active else 0, user_id))
    conn.commit()

def set_user_role(user_id, role: str):
    conn = get_conn()
    c = conn.cursor()
    c.execute("UPDATE users SET role=? WHERE id=?", (role, user_id))
    conn.commit()

def update_user_password(user_id, new_password_hash):
    conn = get_conn()
    c = conn.cursor()
    c.execute("UPDATE users SET password_hash=? WHERE id=?", (new_password_hash, user_id))
    conn.commit()

def delete_user(user_id):
    conn = get_conn()
    c = conn.cursor()
    c.execute("DELETE FROM users WHERE id=?", (user_id,))
    conn.commit()

# --- Conversations & Messages ---
def create_conversation(user_id, title="New Chat"):
    conn = get_conn()
    c = conn.cursor()
    c.execute("INSERT INTO conversations (user_id, title, created_at) VALUES (?,?,?)",
              (user_id, title, datetime.utcnow().isoformat()))
    conn.commit()
    return c.lastrowid

def list_conversations(user_id):
    conn = get_conn()
    c = conn.cursor()
    c.execute("SELECT id, title, created_at FROM conversations WHERE user_id=? ORDER BY id DESC", (user_id,))
    rows = c.fetchall()
    return [{"id": r[0], "title": r[1], "created_at": r[2]} for r in rows]

def rename_conversation(conversation_id, title):
    conn = get_conn()
    c = conn.cursor()
    c.execute("UPDATE conversations SET title=? WHERE id=?", (title, conversation_id))
    conn.commit()

def delete_conversation(conversation_id):
    conn = get_conn()
    c = conn.cursor()
    c.execute("DELETE FROM messages WHERE conversation_id=?", (conversation_id,))
    c.execute("DELETE FROM conversations WHERE id=?", (conversation_id,))
    conn.commit()

def add_message(conversation_id, role, content, attachments=None):
    conn = get_conn()
    c = conn.cursor()
    c.execute("INSERT INTO messages (conversation_id, role, content, attachments, created_at) VALUES (?,?,?,?,?)",
              (conversation_id, role, content, attachments or "[]", datetime.utcnow().isoformat()))
    conn.commit()
    return c.lastrowid

def get_messages(conversation_id):
    conn = get_conn()
    c = conn.cursor()
    c.execute("SELECT role, content, attachments, created_at FROM messages WHERE conversation_id=? ORDER BY id ASC",
              (conversation_id,))
    rows = c.fetchall()
    msgs = []
    for r in rows:
        msgs.append({"role": r[0], "content": r[1], "attachments": r[2], "created_at": r[3]})
    return msgs