database.py

import sqlite3
from datetime import datetime
from contextlib import contextmanager
from typing import Optional
from dataclasses import dataclass
from enum import Enum
import os

from config import DATABASE_PATH


class QueuePhase(Enum):
    REGISTRATION = "registration"
    ACTIVE = "active"
    ADDITIONAL_REGISTRATION = "additional"
    CLOSED = "closed"


class ParticipantStatus(Enum):
    PENDING = "pending"
    MAIN = "main"
    RESERVE = "reserve"
    WAITING = "waiting"
    DECLINED = "declined"
    PROMOTED = "promoted"


@dataclass
class Queue:
    id: str
    name: str
    chat_id: int
    message_id: Optional[int]
    main_slots: int
    reserve_slots: int
    phase: QueuePhase
    created_at: datetime
    closes_at: datetime


@dataclass
class Participant:
    id: int
    queue_id: str
    user_id: int
    username: Optional[str]
    full_name: str
    position: Optional[int]
    status: ParticipantStatus
    joined_at: datetime


class Database:
    def __init__(self, db_path: str = DATABASE_PATH):
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        self.db_path = db_path
        self._init_db()
    
    @contextmanager
    def _get_conn(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        finally:
            conn.close()
    
    def _init_db(self):
        with self._get_conn() as conn:
            conn.executescript("""
                CREATE TABLE IF NOT EXISTS queues (
                    id TEXT PRIMARY KEY,
                    name TEXT NOT NULL,
                    chat_id INTEGER NOT NULL,
                    message_id INTEGER,
                    main_slots INTEGER NOT NULL,
                    reserve_slots INTEGER NOT NULL,
                    phase TEXT NOT NULL DEFAULT 'registration',
                    created_at TEXT NOT NULL,
                    closes_at TEXT NOT NULL
                );
                
                CREATE TABLE IF NOT EXISTS participants (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    queue_id TEXT NOT NULL,
                    user_id INTEGER NOT NULL,
                    username TEXT,
                    full_name TEXT NOT NULL,
                    position INTEGER,
                    status TEXT NOT NULL DEFAULT 'pending',
                    joined_at TEXT NOT NULL,
                    FOREIGN KEY (queue_id) REFERENCES queues(id),
                    UNIQUE(queue_id, user_id)
                );
                
                CREATE INDEX IF NOT EXISTS idx_queue_chat ON queues(chat_id);
                CREATE INDEX IF NOT EXISTS idx_participant_queue ON participants(queue_id);
            """)
    
    def create_queue(self, queue: Queue) -> None:
        with self._get_conn() as conn:
            conn.execute("""
                INSERT INTO queues (id, name, chat_id, message_id, main_slots, reserve_slots, 
                                   phase, created_at, closes_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                queue.id, queue.name, queue.chat_id, queue.message_id,
                queue.main_slots, queue.reserve_slots, queue.phase.value,
                queue.created_at.isoformat(), queue.closes_at.isoformat()
            ))
    
    def get_queue(self, queue_id: str) -> Optional[Queue]:
        with self._get_conn() as conn:
            row = conn.execute("SELECT * FROM queues WHERE id = ?", (queue_id,)).fetchone()
            if row:
                return self._row_to_queue(row)
        return None
    
    def get_active_queues(self, chat_id: int) -> list[Queue]:
        with self._get_conn() as conn:
            rows = conn.execute(
                "SELECT * FROM queues WHERE chat_id = ? AND phase != 'closed'",
                (chat_id,)
            ).fetchall()
            return [self._row_to_queue(r) for r in rows]
    
    def get_registration_queues(self) -> list[Queue]:
        with self._get_conn() as conn:
            rows = conn.execute(
                "SELECT * FROM queues WHERE phase IN ('registration', 'additional')"
            ).fetchall()
            return [self._row_to_queue(r) for r in rows]
    
    def update_queue(self, queue: Queue) -> None:
        with self._get_conn() as conn:
            conn.execute("""
                UPDATE queues SET message_id = ?, phase = ?, closes_at = ?
                WHERE id = ?
            """, (queue.message_id, queue.phase.value, queue.closes_at.isoformat(), queue.id))
    
    def find_queue_by_name(self, chat_id: int, name: str) -> Optional[Queue]:
        with self._get_conn() as conn:
            row = conn.execute(
                "SELECT * FROM queues WHERE chat_id = ? AND LOWER(name) = LOWER(?) AND phase != 'closed'",
                (chat_id, name)
            ).fetchone()
            if row:
                return self._row_to_queue(row)
        return None
    
    def _row_to_queue(self, row) -> Queue:
        return Queue(
            id=row["id"],
            name=row["name"],
            chat_id=row["chat_id"],
            message_id=row["message_id"],
            main_slots=row["main_slots"],
            reserve_slots=row["reserve_slots"],
            phase=QueuePhase(row["phase"]),
            created_at=datetime.fromisoformat(row["created_at"]),
            closes_at=datetime.fromisoformat(row["closes_at"])
        )
    
    def add_participant(self, p: Participant) -> None:
        with self._get_conn() as conn:
            conn.execute("""
                INSERT OR REPLACE INTO participants 
                (queue_id, user_id, username, full_name, position, status, joined_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (
                p.queue_id, p.user_id, p.username, p.full_name,
                p.position, p.status.value, p.joined_at.isoformat()
            ))
    
    def remove_participant(self, queue_id: str, user_id: int) -> bool:
        with self._get_conn() as conn:
            cursor = conn.execute(
                "DELETE FROM participants WHERE queue_id = ? AND user_id = ?",
                (queue_id, user_id)
            )
            return cursor.rowcount > 0
    
    def get_participant(self, queue_id: str, user_id: int) -> Optional[Participant]:
        with self._get_conn() as conn:
            row = conn.execute(
                "SELECT * FROM participants WHERE queue_id = ? AND user_id = ?",
                (queue_id, user_id)
            ).fetchone()
            if row:
                return self._row_to_participant(row)
        return None
    
    def get_participants(self, queue_id: str, status: Optional[ParticipantStatus] = None) -> list[Participant]:
        with self._get_conn() as conn:
            if status:
                rows = conn.execute(
                    "SELECT * FROM participants WHERE queue_id = ? AND status = ? ORDER BY position",
                    (queue_id, status.value)
                ).fetchall()
            else:
                rows = conn.execute(
                    "SELECT * FROM participants WHERE queue_id = ? ORDER BY position, joined_at",
                    (queue_id,)
                ).fetchall()
            return [self._row_to_participant(r) for r in rows]
    
    def get_pending_participants(self, queue_id: str) -> list[Participant]:
        return self.get_participants(queue_id, ParticipantStatus.PENDING)
    
    def get_main_participants(self, queue_id: str) -> list[Participant]:
        with self._get_conn() as conn:
            rows = conn.execute(
                """SELECT * FROM participants 
                   WHERE queue_id = ? AND status IN ('main', 'promoted') 
                   ORDER BY position""",
                (queue_id,)
            ).fetchall()
            return [self._row_to_participant(r) for r in rows]
    
    def get_reserve_participants(self, queue_id: str) -> list[Participant]:
        return self.get_participants(queue_id, ParticipantStatus.RESERVE)
    
    def update_participant_status(self, queue_id: str, user_id: int, 
                                   status: ParticipantStatus, position: Optional[int] = None) -> None:
        with self._get_conn() as conn:
            if position is not None:
                conn.execute(
                    "UPDATE participants SET status = ?, position = ? WHERE queue_id = ? AND user_id = ?",
                    (status.value, position, queue_id, user_id)
                )
            else:
                conn.execute(
                    "UPDATE participants SET status = ? WHERE queue_id = ? AND user_id = ?",
                    (status.value, queue_id, user_id)
                )
    
    def count_participants(self, queue_id: str, exclude_declined: bool = True) -> int:
        with self._get_conn() as conn:
            if exclude_declined:
                row = conn.execute(
                    "SELECT COUNT(*) as cnt FROM participants WHERE queue_id = ? AND status != 'declined'",
                    (queue_id,)
                ).fetchone()
            else:
                row = conn.execute(
                    "SELECT COUNT(*) as cnt FROM participants WHERE queue_id = ?",
                    (queue_id,)
                ).fetchone()
            return row["cnt"]
    
    def get_first_reserve(self, queue_id: str) -> Optional[Participant]:
        with self._get_conn() as conn:
            row = conn.execute(
                """SELECT * FROM participants 
                   WHERE queue_id = ? AND status = 'reserve' 
                   ORDER BY position LIMIT 1""",
                (queue_id,)
            ).fetchone()
            if row:
                return self._row_to_participant(row)
        return None
    
    def _row_to_participant(self, row) -> Participant:
        return Participant(
            id=row["id"],
            queue_id=row["queue_id"],
            user_id=row["user_id"],
            username=row["username"],
            full_name=row["full_name"],
            position=row["position"],
            status=ParticipantStatus(row["status"]),
            joined_at=datetime.fromisoformat(row["joined_at"])
        )


db = Database()