database.py

"""
Asynchronous SQLite database interface for Parcel Tracker Bot.
Manages users, parcels, tracking events, and monitoring logs.
"""

import json
import logging
import aiosqlite
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any, Tuple
import config

logger = logging.getLogger(__name__)


class Database:
    """Async SQLite Database Manager."""

    def __init__(self, db_path: str = config.DB_PATH):
        self.db_path = db_path

    async def init_db(self) -> None:
        """Initialize database schema and tables."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("PRAGMA journal_mode=WAL;")
            await db.execute("PRAGMA foreign_keys=ON;")

            # Users table
            await db.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    user_id INTEGER PRIMARY KEY,
                    username TEXT,
                    first_name TEXT,
                    language_code TEXT DEFAULT 'ru',
                    is_admin INTEGER DEFAULT 0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_activity TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)

            # Parcels table
            await db.execute("""
                CREATE TABLE IF NOT EXISTS parcels (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER NOT NULL,
                    tracking_number TEXT NOT NULL,
                    title TEXT NOT NULL,
                    carrier TEXT DEFAULT 'Unknown Carrier',
                    status TEXT DEFAULT 'Pending',
                    last_event TEXT DEFAULT '',
                    location TEXT DEFAULT '',
                    event_date TEXT DEFAULT '',
                    estimated_delivery TEXT DEFAULT '',
                    history_json TEXT DEFAULT '[]',
                    is_archived INTEGER DEFAULT 0,
                    notify_enabled INTEGER DEFAULT 1,
                    last_check_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(user_id, tracking_number),
                    FOREIGN KEY(user_id) REFERENCES users(user_id) ON DELETE CASCADE
                );
            """)

            # Check logs table for monitoring & admin stats
            await db.execute("""
                CREATE TABLE IF NOT EXISTS check_logs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER,
                    tracking_number TEXT,
                    is_success INTEGER,
                    error_message TEXT,
                    checked_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)

            # Indexes
            await db.execute("CREATE INDEX IF NOT EXISTS idx_parcels_user ON parcels(user_id);")
            await db.execute("CREATE INDEX IF NOT EXISTS idx_parcels_track ON parcels(tracking_number);")
            await db.execute("CREATE INDEX IF NOT EXISTS idx_parcels_active ON parcels(is_archived, notify_enabled);")
            await db.commit()
            logger.info("Database initialized successfully.")

    # ---------------- USER METHODS ----------------

    async def get_or_create_user(
        self,
        user_id: int,
        username: Optional[str] = None,
        first_name: Optional[str] = None,
        telegram_lang: Optional[str] = None
    ) -> Dict[str, Any]:
        """Fetch existing user or create a new one."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("SELECT * FROM users WHERE user_id = ?", (user_id,)) as cursor:
                row = await cursor.fetchone()
                if row:
                    await db.execute(
                        "UPDATE users SET last_activity = CURRENT_TIMESTAMP, username = ?, first_name = ? WHERE user_id = ?",
                        (username, first_name, user_id)
                    )
                    await db.commit()
                    return dict(row)

            # Determine initial language
            lang = 'ru'
            if telegram_lang and telegram_lang.startswith('en'):
                lang = 'en'

            is_admin = 1 if user_id in config.ADMIN_IDS else 0

            await db.execute("""
                INSERT INTO users (user_id, username, first_name, language_code, is_admin)
                VALUES (?, ?, ?, ?, ?)
            """, (user_id, username, first_name, lang, is_admin))
            await db.commit()

            return {
                'user_id': user_id,
                'username': username,
                'first_name': first_name,
                'language_code': lang,
                'is_admin': is_admin
            }

    async def get_user_language(self, user_id: int) -> str:
        """Get user's preferred language."""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("SELECT language_code FROM users WHERE user_id = ?", (user_id,)) as cursor:
                row = await cursor.fetchone()
                return row[0] if row and row[0] else 'ru'

    async def set_user_language(self, user_id: int, lang_code: str) -> None:
        """Update user's language code."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("UPDATE users SET language_code = ? WHERE user_id = ?", (lang_code, user_id))
            await db.commit()

    async def get_all_user_ids(self) -> List[int]:
        """Fetch all user IDs (for broadcast)."""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("SELECT user_id FROM users") as cursor:
                rows = await cursor.fetchall()
                return [r[0] for r in rows]

    # ---------------- PARCEL METHODS ----------------

    async def save_or_update_parcel(
        self,
        user_id: int,
        tracking_number: str,
        title: Optional[str] = None,
        carrier: Optional[str] = None,
        status: Optional[str] = None,
        last_event: Optional[str] = None,
        location: Optional[str] = None,
        event_date: Optional[str] = None,
        estimated_delivery: Optional[str] = None,
        history: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """Save a new parcel or update tracking data for an existing one."""
        tracking_number = tracking_number.strip().upper()
        history_str = json.dumps(history or [], ensure_ascii=False)

        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute(
                "SELECT * FROM parcels WHERE user_id = ? AND tracking_number = ?",
                (user_id, tracking_number)
            ) as cursor:
                existing = await cursor.fetchone()

            if existing:
                new_title = title if title else existing['title']
                await db.execute("""
                    UPDATE parcels SET
                        title = ?,
                        carrier = COALESCE(?, carrier),
                        status = COALESCE(?, status),
                        last_event = COALESCE(?, last_event),
                        location = COALESCE(?, location),
                        event_date = COALESCE(?, event_date),
                        estimated_delivery = COALESCE(?, estimated_delivery),
                        history_json = CASE WHEN ? != '[]' THEN ? ELSE history_json END,
                        last_check_at = CURRENT_TIMESTAMP,
                        updated_at = CURRENT_TIMESTAMP
                    WHERE user_id = ? AND tracking_number = ?
                """, (
                    new_title, carrier, status, last_event, location, event_date,
                    estimated_delivery, history_str, history_str, user_id, tracking_number
                ))
            else:
                default_title = title if title else f"Посылка {tracking_number}"
                await db.execute("""
                    INSERT INTO parcels (
                        user_id, tracking_number, title, carrier, status,
                        last_event, location, event_date, estimated_delivery,
                        history_json, is_archived, notify_enabled
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, 1)
                """, (
                    user_id, tracking_number, default_title,
                    carrier or 'Unknown Carrier', status or 'In transit',
                    last_event or '', location or '', event_date or '',
                    estimated_delivery or '', history_str
                ))

            await db.commit()

            async with db.execute(
                "SELECT * FROM parcels WHERE user_id = ? AND tracking_number = ?",
                (user_id, tracking_number)
            ) as cursor:
                return dict(await cursor.fetchone())

    async def get_parcel_by_id(self, parcel_id: int) -> Optional[Dict[str, Any]]:
        """Retrieve parcel by primary key."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("SELECT * FROM parcels WHERE id = ?", (parcel_id,)) as cursor:
                row = await cursor.fetchone()
                return dict(row) if row else None

    async def get_parcel_by_user_and_track(self, user_id: int, tracking_number: str) -> Optional[Dict[str, Any]]:
        """Retrieve parcel by user ID and tracking number."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute(
                "SELECT * FROM parcels WHERE user_id = ? AND tracking_number = ?",
                (user_id, tracking_number.strip().upper())
            ) as cursor:
                row = await cursor.fetchone()
                return dict(row) if row else None

    async def get_user_parcels(
        self,
        user_id: int,
        is_archived: bool = False,
        limit: int = 5,
        offset: int = 0
    ) -> Tuple[List[Dict[str, Any]], int]:
        """Fetch user parcels with pagination and total count."""
        archived_val = 1 if is_archived else 0
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            # Count total
            async with db.execute(
                "SELECT COUNT(*) FROM parcels WHERE user_id = ? AND is_archived = ?",
                (user_id, archived_val)
            ) as cursor:
                total_count = (await cursor.fetchone())[0]

            # Fetch page
            async with db.execute(
                """
                SELECT * FROM parcels
                WHERE user_id = ? AND is_archived = ?
                ORDER BY updated_at DESC, id DESC
                LIMIT ? OFFSET ?
                """,
                (user_id, archived_val, limit, offset)
            ) as cursor:
                rows = await cursor.fetchall()
                return [dict(r) for r in rows], total_count

    async def update_parcel_title(self, parcel_id: int, user_id: int, new_title: str) -> bool:
        """Rename a parcel."""
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute(
                "UPDATE parcels SET title = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ? AND user_id = ?",
                (new_title.strip(), parcel_id, user_id)
            )
            await db.commit()
            return cursor.rowcount > 0

    async def toggle_parcel_archive(self, parcel_id: int, user_id: int) -> Optional[bool]:
        """Toggle archive status of a parcel."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("SELECT is_archived FROM parcels WHERE id = ? AND user_id = ?", (parcel_id, user_id)) as cursor:
                row = await cursor.fetchone()
                if not row:
                    return None
                new_state = 0 if row['is_archived'] else 1

            await db.execute(
                "UPDATE parcels SET is_archived = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
                (new_state, parcel_id)
            )
            await db.commit()
            return bool(new_state)

    async def toggle_parcel_notify(self, parcel_id: int, user_id: int) -> Optional[bool]:
        """Toggle notification status of a parcel."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("SELECT notify_enabled FROM parcels WHERE id = ? AND user_id = ?", (parcel_id, user_id)) as cursor:
                row = await cursor.fetchone()
                if not row:
                    return None
                new_state = 0 if row['notify_enabled'] else 1

            await db.execute(
                "UPDATE parcels SET notify_enabled = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?",
                (new_state, parcel_id)
            )
            await db.commit()
            return bool(new_state)

    async def delete_parcel(self, parcel_id: int, user_id: int) -> bool:
        """Delete a parcel."""
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute("DELETE FROM parcels WHERE id = ? AND user_id = ?", (parcel_id, user_id))
            await db.commit()
            return cursor.rowcount > 0

    async def get_active_parcels_for_background_check(self) -> List[Dict[str, Any]]:
        """Fetch all non-archived parcels with notifications enabled for background updater."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("""
                SELECT p.*, u.language_code
                FROM parcels p
                JOIN users u ON p.user_id = u.user_id
                WHERE p.is_archived = 0 AND p.notify_enabled = 1
                ORDER BY p.last_check_at ASC
            """) as cursor:
                rows = await cursor.fetchall()
                return [dict(r) for r in rows]

    # ---------------- LOGS & ADMIN STATS ----------------

    async def log_check(
        self,
        user_id: int,
        tracking_number: str,
        is_success: bool,
        error_message: Optional[str] = None
    ) -> None:
        """Log a tracking check."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                INSERT INTO check_logs (user_id, tracking_number, is_success, error_message)
                VALUES (?, ?, ?, ?)
            """, (user_id, tracking_number, 1 if is_success else 0, error_message or ''))
            await db.commit()

    async def get_admin_stats(self) -> Dict[str, Any]:
        """Fetch administrative summary statistics."""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("SELECT COUNT(*) FROM users") as c:
                total_users = (await c.fetchone())[0]

            async with db.execute("SELECT COUNT(*) FROM parcels") as c:
                total_parcels = (await c.fetchone())[0]

            async with db.execute("SELECT COUNT(*) FROM parcels WHERE is_archived = 0") as c:
                active_parcels = (await c.fetchone())[0]

            async with db.execute("SELECT COUNT(*) FROM parcels WHERE is_archived = 1") as c:
                archived_parcels = (await c.fetchone())[0]

            # 24h logs
            one_day_ago = (datetime.utcnow() - timedelta(days=1)).strftime('%Y-%m-%d %H:%M:%S')
            async with db.execute(
                "SELECT COUNT(*) FROM check_logs WHERE checked_at >= ?", (one_day_ago,)
            ) as c:
                checks_24h = (await c.fetchone())[0]

            async with db.execute(
                "SELECT COUNT(*) FROM check_logs WHERE checked_at >= ? AND is_success = 0", (one_day_ago,)
            ) as c:
                errors_24h = (await c.fetchone())[0]

            return {
                'total_users': total_users,
                'total_parcels': total_parcels,
                'active_parcels': active_parcels,
                'archived_parcels': archived_parcels,
                'checks_24h': checks_24h,
                'errors_24h': errors_24h,
                'interval_min': config.CHECK_INTERVAL_SECONDS // 60
            }


# Singleton database instance
db = Database()