src/firestore_lottery.py

from google.cloud import firestore
from google.cloud.firestore_v1.async_client import AsyncClient
from google.cloud.firestore_v1 import FieldFilter
from datetime import datetime
from typing import Dict, List, Optional
from src.logger import logger
import asyncio
from tenacity import retry, stop_after_attempt, wait_fixed
import traceback
from config.settings import Settings

class FirestoreLottery:
    def __init__(self, credentials: Dict):
        logger.debug("[FirestoreLottery] Initializing with credentials")
        if not credentials or not isinstance(credentials, dict):
            logger.error(f"[FirestoreLottery] Invalid or missing FIRESTORE_CREDENTIALS: {credentials}")
            self.async_db = None
            return
        try:
            logger.debug(f"[FirestoreLottery] Credentials keys: {list(credentials.keys())}")
            from google.oauth2.service_account import Credentials
            creds = Credentials.from_service_account_info(credentials)
            self.async_db = AsyncClient(credentials=creds)
            self.lottery_results = self.async_db.collection("lottery-results")
            logger.info("[FirestoreLottery] Initialized Firestore client successfully")
        except Exception as e:
            logger.error(f"[FirestoreLottery] Initialization error: {str(e)}\n{traceback.format_exc()}")
            self.async_db = None

    async def parse_date(self, date_str: str) -> Optional[tuple]:
        formats = [
            "%d-%m-%Y", "%Y-%m-%d", "%d/%m/%Y", "%Y/%m/%d",
            "%d-%m-%y", "%y-%m-%d", "%d/%m/%y", "%y/%m/%d",
            "%d.%m.%Y", "%d-%m-%Y",
            "%d/%m/%Y", "%Y.%m.%d", "%m/%d/%Y", "%m-%d-%Y",
            "%d%m%Y", "%Y%m%d"
        ]
        for fmt in formats:
            try:
                parsed = datetime.strptime(date_str, fmt)
                return parsed.strftime("%d-%m-%Y"), parsed
            except ValueError:
                continue
        logger.warning(f"[FirestoreLottery] Invalid date format: {date_str}")
        return None

    def _map_dai_name(self, dai: str) -> str:
        dai_clean = dai.lower().replace(" ", "").replace("chủnhật", "").replace("chu", "")
        for dai_code, config in Settings.DAI_CONFIG.items():
            config_name_clean = config["name"].lower().replace(" ", "")
            if config_name_clean in dai_clean or dai_clean in config_name_clean:
                return dai_code
        return dai.lower()

    @retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
    async def get_lottery_by_dai(self, dai: str, limit: int = 500, select_fields: Optional[List[str]] = None) -> List[Dict]:
        if not self.async_db:
            logger.error("[get_lottery_by_dai] FirestoreDB not initialized")
            return []
        try:
            async with asyncio.timeout(Settings.QUERY_TIMEOUT or 60):
                dai_normalized = self._map_dai_name(dai)
                logger.debug(f"[get_lottery_by_dai] Querying for dai={dai_normalized}, original={dai}, limit={limit}, select_fields={select_fields}")
                query = self.lottery_results.where(
                    filter=firestore.FieldFilter("dai", "==", dai_normalized)
                )
                if select_fields:
                    query = query.select(select_fields)
                query = query.order_by("ngay", direction=firestore.Query.DESCENDING).limit(limit)
                results = []
                async for doc in query.stream():
                    data = doc.to_dict()
                    logger.debug(f"[get_lottery_by_dai] Raw document: {data}")
                    ngay = data.get('ngay')
                    if ngay:
                        parsed = await self.parse_date(ngay)
                        if parsed:
                            data['ngay'] = parsed[0]
                            results.append(data)
                        else:
                            logger.warning(f"[get_lottery_by_dai] Skip invalid date: ngay={ngay}, dai={dai_normalized}")
                    else:
                        logger.warning(f"[get_lottery_by_dai] Skip document without ngay: {data}")
                logger.info(f"[get_lottery_by_dai] Retrieved {len(results)} records for dai={dai_normalized}")
                if not results:
                    logger.debug(f"[get_lottery_by_dai] No results for {dai_normalized}, trying all documents")
                    query_all = self.lottery_results.limit(10)
                    async for doc in query_all.stream():
                        logger.debug(f"[get_lottery_by_dai] Available document: {doc.to_dict()}")
                return results
        except asyncio.TimeoutError:
            logger.error(f"[get_lottery_by_dai] Timeout: dai={dai}")
            return []
        except Exception as e:
            logger.error(f"[get_lottery_by_dai] Error: dai={dai}, error={str(e)}\n{traceback.format_exc()}")
            return []