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()