database.py

"""
AI 垃圾分类助手 - 数据库模块
使用 SQLite 记录用户分类历史和环保积分
"""

import sqlite3
from datetime import date


class Database:
    def __init__(self, db_path="garbage_assistant.db"):
        self.db_path = db_path
        self.init_database()

    def _get_conn(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        return conn

    def init_database(self):
        conn = self._get_conn()
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                total_points INTEGER DEFAULT 0,
                total_classifications INTEGER DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                image_path TEXT,
                predicted_class TEXT NOT NULL,
                confidence REAL,
                points INTEGER DEFAULT 10,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users(id)
            );

            CREATE TABLE IF NOT EXISTS daily_stats (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                date TEXT NOT NULL,
                count INTEGER DEFAULT 0,
                points INTEGER DEFAULT 0,
                FOREIGN KEY (user_id) REFERENCES users(id),
                UNIQUE(user_id, date)
            );
        """)
        conn.commit()
        conn.close()

    def register_user(self, username):
        conn = self._get_conn()
        try:
            cursor = conn.cursor()
            cursor.execute("INSERT INTO users (username) VALUES (?)", (username,))
            conn.commit()
            user_id = cursor.lastrowid
            print(f"✓ 用户 '{username}' 注册成功 (ID: {user_id})")
            return user_id
        except sqlite3.IntegrityError:
            user_id = conn.execute(
                "SELECT id FROM users WHERE username = ?", (username,)
            ).fetchone()["id"]
            print(f"ℹ 用户 '{username}' 已存在 (ID: {user_id})")
            return user_id
        finally:
            conn.close()

    def get_user(self, username):
        conn = self._get_conn()
        user = conn.execute(
            "SELECT * FROM users WHERE username = ?", (username,)
        ).fetchone()
        conn.close()
        return dict(user) if user else None

    def add_record(self, user_id, predicted_class, confidence, image_path=None):
        today = date.today().isoformat()
        points = max(5, min(20, int(confidence * 20)))

        conn = self._get_conn()
        conn.execute(
            "INSERT INTO records (user_id, image_path, predicted_class, confidence, points) VALUES (?, ?, ?, ?, ?)",
            (user_id, image_path, predicted_class, confidence, points),
        )
        conn.execute(
            "UPDATE users SET total_points = total_points + ?, total_classifications = total_classifications + 1 WHERE id = ?",
            (points, user_id),
        )
        conn.execute(
            "INSERT INTO daily_stats (user_id, date, count, points) VALUES (?, ?, 1, ?) "
            "ON CONFLICT(user_id, date) DO UPDATE SET count = count + 1, points = points + ?",
            (user_id, today, points, points),
        )
        conn.commit()
        conn.close()
        return points

    def get_user_stats(self, user_id):
        conn = self._get_conn()
        user = conn.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
        if not user:
            conn.close()
            return None

        stats = dict(user)
        stats["recent_records"] = [
            dict(r) for r in conn.execute(
                "SELECT predicted_class, confidence, points, created_at FROM records WHERE user_id = ? ORDER BY created_at DESC LIMIT 10",
                (user_id,),
            ).fetchall()
        ]
        stats["class_distribution"] = [
            dict(r) for r in conn.execute(
                "SELECT predicted_class, COUNT(*) as count FROM records WHERE user_id = ? GROUP BY predicted_class ORDER BY count DESC",
                (user_id,),
            ).fetchall()
        ]

        today = date.today().isoformat()
        today_row = conn.execute(
            "SELECT * FROM daily_stats WHERE user_id = ? AND date = ?", (user_id, today)
        ).fetchone()
        stats["today"] = dict(today_row) if today_row else {"count": 0, "points": 0}

        stats["weekly"] = [
            dict(r) for r in conn.execute(
                "SELECT date, count, points FROM daily_stats WHERE user_id = ? AND date >= date('now', '-7 days') ORDER BY date DESC",
                (user_id,),
            ).fetchall()
        ]

        conn.close()
        return stats

    def get_leaderboard(self, limit=10):
        conn = self._get_conn()
        leaders = conn.execute(
            "SELECT username, total_points, total_classifications FROM users ORDER BY total_points DESC LIMIT ?",
            (limit,),
        ).fetchall()
        conn.close()
        return [dict(r) for r in leaders]