scripts/database.py

import sqlite3
import os
import secrets
import hashlib
from typing import Optional, Dict, List

class DatabaseManager:
    def __init__(self, db_path=None):
        self.db_path = db_path or os.environ.get('DATABASE_PATH', 'tmc_customer_service.db')
        os.makedirs(os.path.dirname(self.db_path) or '.', exist_ok=True)
        self.init_database()

    def get_connection(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn

    def init_database(self):
        with self.get_connection() as conn:
            c = conn.cursor()
            c.execute('''CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                email TEXT UNIQUE NOT NULL,
                first_name TEXT NOT NULL,
                last_name TEXT NOT NULL,
                password_hash TEXT NOT NULL,
                salt TEXT NOT NULL,
                phone TEXT,
                company TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                is_active BOOLEAN DEFAULT 1,
                role TEXT DEFAULT 'user'
            )''')
            c.execute('''CREATE TABLE IF NOT EXISTS sessions (
                id TEXT PRIMARY KEY,
                user_id INTEGER,
                ip_address TEXT,
                user_agent TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                expires_at TIMESTAMP NOT NULL,
                is_active BOOLEAN DEFAULT 1,
                FOREIGN KEY(user_id) REFERENCES users(id)
            )''')
            c.execute('''CREATE TABLE IF NOT EXISTS conversations (
                id TEXT PRIMARY KEY,
                user_id INTEGER,
                session_id TEXT,
                title TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                is_active BOOLEAN DEFAULT 1,
                FOREIGN KEY(user_id) REFERENCES users(id)
            )''')
            c.execute('''CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                conversation_id TEXT NOT NULL,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                model_used TEXT,
                response_time_ms INTEGER,
                FOREIGN KEY(conversation_id) REFERENCES conversations(id)
            )''')
            c.execute('''CREATE TABLE IF NOT EXISTS support_tickets (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ticket_number TEXT UNIQUE NOT NULL,
                user_id INTEGER NOT NULL,
                conversation_id TEXT,
                subject TEXT NOT NULL,
                description TEXT NOT NULL,
                category TEXT NOT NULL,
                priority TEXT DEFAULT 'medium',
                status TEXT DEFAULT 'open',
                assigned_agent TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                resolved_at TIMESTAMP,
                FOREIGN KEY(user_id) REFERENCES users(id)
            )''')
            c.execute('''CREATE TABLE IF NOT EXISTS ticket_updates (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ticket_id INTEGER NOT NULL,
                user_id INTEGER,
                update_type TEXT DEFAULT 'note',
                message TEXT NOT NULL,
                is_internal BOOLEAN DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY(ticket_id) REFERENCES support_tickets(id)
            )''')
            c.execute('''CREATE TABLE IF NOT EXISTS ticket_categories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT UNIQUE NOT NULL,
                description TEXT,
                default_priority TEXT DEFAULT 'medium',
                is_active BOOLEAN DEFAULT 1
            )''')
            conn.commit()
            c.execute("INSERT OR IGNORE INTO ticket_categories (name, description) VALUES ('General', 'General inquiries')")
            conn.commit()

    def hash_password(self, password: str):
        salt = secrets.token_hex(32)
        ph = hashlib.pbkdf2_hmac('sha256', password.encode(), salt.encode(), 100000).hex()
        return ph, salt

    def verify_password(self, password, phash, salt):
        return hashlib.pbkdf2_hmac('sha256', password.encode(), salt.encode(), 100000).hex() == phash

    def create_user(self, email, first_name, last_name, password, phone=None, company=None) -> Optional[int]:
        try:
            phash, salt = self.hash_password(password)
            with self.get_connection() as conn:
                cur = conn.cursor()
                cur.execute('''INSERT INTO users (email, first_name, last_name, password_hash, salt, phone, company)
                               VALUES (?,?,?,?,?,?,?)''',
                            (email, first_name, last_name, phash, salt, phone, company))
                uid = cur.lastrowid
                conn.commit()
                return uid
        except sqlite3.IntegrityError:
            return None

    def authenticate_user(self, email, password) -> Optional[Dict]:
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('SELECT id, email, first_name, last_name, password_hash, salt, is_active FROM users WHERE email = ?', (email,))
            row = cur.fetchone()
            if row and row['is_active'] and self.verify_password(password, row['password_hash'], row['salt']):
                return dict(row)
            return None

    def create_session(self, user_id, ip, ua, hours=24) -> str:
        sid = secrets.token_urlsafe(32)
        with self.get_connection() as conn:
            conn.execute('INSERT INTO sessions (id, user_id, ip_address, user_agent, expires_at) VALUES (?,?,?,?, datetime("now", "+? hours"))',
                         (sid, user_id, ip, ua[:255], hours))
            conn.commit()
        return sid

    def get_user_by_session(self, session_id) -> Optional[Dict]:
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('''SELECT u.id, u.email, u.first_name, u.last_name, u.role
                           FROM users u JOIN sessions s ON u.id = s.user_id
                           WHERE s.id = ? AND s.is_active = 1 AND s.expires_at > CURRENT_TIMESTAMP''', (session_id,))
            row = cur.fetchone()
            return dict(row) if row else None

    def get_user_role(self, user_id) -> str:
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('SELECT role FROM users WHERE id = ?', (user_id,))
            row = cur.fetchone()
            return row['role'] if row else 'user'

    def create_conversation(self, user_id=None, session_id=None, title=None) -> str:
        cid = secrets.token_urlsafe(16)
        with self.get_connection() as conn:
            conn.execute('INSERT INTO conversations (id, user_id, session_id, title) VALUES (?,?,?,?)',
                         (cid, user_id, session_id, title))
            conn.commit()
        return cid

    def add_message(self, conversation_id, role, content, model_used=None, response_time_ms=None):
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('''INSERT INTO messages (conversation_id, role, content, model_used, response_time_ms)
                           VALUES (?,?,?,?,?)''', (conversation_id, role, content, model_used, response_time_ms))
            conn.execute('UPDATE conversations SET updated_at = CURRENT_TIMESTAMP WHERE id = ?', (conversation_id,))
            conn.commit()
            return cur.lastrowid

    def get_conversation_history(self, conversation_id, limit=50) -> List[Dict]:
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('SELECT is_active FROM conversations WHERE id = ?', (conversation_id,))
            if not cur.fetchone():
                return []
            cur.execute('''SELECT role, content, timestamp, model_used FROM messages
                           WHERE conversation_id = ? ORDER BY timestamp LIMIT ?''', (conversation_id, limit))
            return [dict(row) for row in cur.fetchall()]

    def create_support_ticket(self, user_id, subject, description, category, conversation_id=None, priority='medium') -> str:
        import random, string
        tn = 'TMC-' + ''.join(random.choices(string.digits, k=6))
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('''INSERT INTO support_tickets (ticket_number, user_id, conversation_id, subject, description, category, priority)
                           VALUES (?,?,?,?,?,?,?)''', (tn, user_id, conversation_id, subject, description, category, priority))
            conn.commit()
            return tn

    def get_ticket_by_number(self, ticket_number) -> Optional[Dict]:
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('SELECT * FROM support_tickets WHERE ticket_number = ?', (ticket_number,))
            row = cur.fetchone()
            return dict(row) if row else None

    def get_ticket_updates(self, ticket_id, include_internal=False) -> List[Dict]:
        with self.get_connection() as conn:
            cur = conn.cursor()
            if include_internal:
                cur.execute('SELECT * FROM ticket_updates WHERE ticket_id = ? ORDER BY created_at', (ticket_id,))
            else:
                cur.execute('SELECT * FROM ticket_updates WHERE ticket_id = ? AND is_internal = 0 ORDER BY created_at', (ticket_id,))
            return [dict(row) for row in cur.fetchall()]

    def add_ticket_update(self, ticket_id, user_id, message, update_type='note', is_internal=False) -> int:
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('''INSERT INTO ticket_updates (ticket_id, user_id, update_type, message, is_internal)
                           VALUES (?,?,?,?,?)''', (ticket_id, user_id, update_type, message, is_internal))
            conn.execute('UPDATE support_tickets SET updated_at = CURRENT_TIMESTAMP WHERE id = ?', (ticket_id,))
            conn.commit()
            return cur.lastrowid

    def get_user_tickets(self, user_id, limit=20) -> List[Dict]:
        with self.get_connection() as conn:
            cur = conn.cursor()
            cur.execute('SELECT * FROM support_tickets WHERE user_id = ? ORDER BY created_at DESC LIMIT ?', (user_id, limit))
            return [dict(row) for row in cur.fetchall()]

    def get_tickets_by_status(self, status, limit=50) -> List[Dict]:
        with self.get_connection() as conn:
            cur = conn.cursor()
            if status:
                cur.execute('SELECT * FROM support_tickets WHERE status = ? ORDER BY created_at DESC LIMIT ?', (status, limit))
            else:
                cur.execute('SELECT * FROM support_tickets ORDER BY created_at DESC LIMIT ?', (limit,))
            return [dict(row) for row in cur.fetchall()]

    # Stubs for additional methods used in original app
    def categorize_ticket_content(self, text):
        return "General"

    def escalate_ticket(self, ticket_id, reason, user_id=None):
        with self.get_connection() as conn:
            conn.execute("UPDATE support_tickets SET priority = 'high' WHERE id = ?", (ticket_id,))
            conn.commit()
            return True

    def check_escalation_needed(self, ticket_id):
        return {"needs_escalation": False}