database.py

import sqlite3
from datetime import datetime, timedelta
from typing import Optional, Dict, List, Any
from enum import Enum
import json
import hashlib

class Platform(str, Enum):
    WHATSAPP = "whatsapp"
    TELEGRAM = "telegram"
    UI = "ui"

class Database:
    def __init__(self, db_path: str = "bot_data.db"):
        self.db_path = db_path
        self.init_db()
    
    def get_connection(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn
    
    def init_db(self):
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                platform TEXT NOT NULL,
                platform_user_id TEXT NOT NULL,
                name TEXT,
                age INTEGER,
                location_str TEXT,
                language_pref TEXT DEFAULT 'en',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(platform, platform_user_id)
            )
        """)
        
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS query_cache (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                category TEXT NOT NULL,
                budget_inr INTEGER NOT NULL,
                location_key TEXT,
                response_text TEXT NOT NULL,
                picks_json TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                ttl_expires_at TIMESTAMP NOT NULL,
                FOREIGN KEY (user_id) REFERENCES users(id)
            )
        """)
        
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS events (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                channel TEXT NOT NULL,
                type TEXT NOT NULL,
                payload_meta TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users(id)
            )
        """)
        
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_users_platform 
            ON users(platform, platform_user_id)
        """)
        
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_cache_lookup 
            ON query_cache(category, budget_inr, location_key, ttl_expires_at)
        """)
        
        conn.commit()
        conn.close()
    
    def get_user(self, platform: Platform, platform_user_id: str) -> Optional[Dict]:
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT * FROM users 
            WHERE platform = ? AND platform_user_id = ?
        """, (platform.value, platform_user_id))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            return dict(row)
        return None
    
    def upsert_user(self, platform: Platform, platform_user_id: str, **fields) -> int:
        conn = self.get_connection()
        cursor = conn.cursor()
        
        existing = self.get_user(platform, platform_user_id)
        
        if existing:
            update_fields = []
            values = []
            for key, value in fields.items():
                if key in ['name', 'age', 'location_str', 'language_pref']:
                    update_fields.append(f"{key} = ?")
                    values.append(value)
            
            if update_fields:
                values.extend([datetime.now(), platform.value, platform_user_id])
                cursor.execute(f"""
                    UPDATE users 
                    SET {', '.join(update_fields)}, updated_at = ?
                    WHERE platform = ? AND platform_user_id = ?
                """, values)
            
            user_id = existing['id']
        else:
            cursor.execute("""
                INSERT INTO users (platform, platform_user_id, name, age, location_str, language_pref)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (
                platform.value, 
                platform_user_id,
                fields.get('name'),
                fields.get('age'),
                fields.get('location_str'),
                fields.get('language_pref', 'en')
            ))
            user_id = cursor.lastrowid
        
        conn.commit()
        conn.close()
        return user_id
    
    def cache_get(self, category: str, budget_inr: int, location_key: str) -> Optional[Dict]:
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT * FROM query_cache
            WHERE category = ? AND budget_inr = ? AND location_key = ?
            AND ttl_expires_at > ?
            ORDER BY created_at DESC
            LIMIT 1
        """, (category, budget_inr, location_key, datetime.now()))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            result = dict(row)
            if result.get('picks_json'):
                result['picks_json'] = json.loads(result['picks_json'])
            return result
        return None
    
    def cache_put(
        self, 
        user_id: int,
        category: str, 
        budget_inr: int, 
        location_key: str,
        response_text: str,
        picks_json: List[Dict],
        ttl_hours: int = 24
    ) -> int:
        conn = self.get_connection()
        cursor = conn.cursor()
        
        ttl_expires_at = datetime.now() + timedelta(hours=ttl_hours)
        
        cursor.execute("""
            INSERT INTO query_cache 
            (user_id, category, budget_inr, location_key, response_text, picks_json, ttl_expires_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        """, (
            user_id,
            category,
            budget_inr,
            location_key,
            response_text,
            json.dumps(picks_json),
            ttl_expires_at
        ))
        
        cache_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return cache_id
    
    def log_event(self, user_id: int, channel: str, event_type: str, payload_meta: Dict = None):
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute("""
            INSERT INTO events (user_id, channel, type, payload_meta)
            VALUES (?, ?, ?, ?)
        """, (user_id, channel, event_type, json.dumps(payload_meta) if payload_meta else None))
        
        conn.commit()
        conn.close()
    
    @staticmethod
    def make_location_key(location_str: Optional[str]) -> str:
        if not location_str:
            return "default"
        return hashlib.md5(location_str.lower().strip().encode()).hexdigest()[:8]

db = Database()