backend/app/services/votes.py

from datetime import datetime, timezone
from typing import Dict, Any, List, Set, Tuple, Optional
import json
import logging
import asyncio
from pathlib import Path
import os
import aiohttp
from huggingface_hub import HfApi
import datasets

from app.services.hf_service import HuggingFaceService
from app.config import HF_TOKEN, API
from app.config.hf_config import HF_ORGANIZATION
from app.core.cache import cache_config
from app.utils.logging import LogFormatter

logger = logging.getLogger(__name__)

class VoteService(HuggingFaceService):
    _instance: Optional['VoteService'] = None
    _initialized = False
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(VoteService, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if not hasattr(self, '_init_done'):
            super().__init__()
            self.votes_file = cache_config.votes_file
            self.votes_to_upload: List[Dict[str, Any]] = []
            self.vote_check_set: Set[Tuple[str, str, str]] = set()
            self._votes_by_model: Dict[str, List[Dict[str, Any]]] = {}
            self._votes_by_user: Dict[str, List[Dict[str, Any]]] = {}
            self._upload_lock = asyncio.Lock()
            self._last_sync = None
            self._sync_interval = 300  # 5 minutes
            self._total_votes = 0
            self._last_vote_timestamp = None
            self._max_retries = 3
            self._retry_delay = 1  # seconds
            self._upload_batch_size = 10
            self.hf_api = HfApi(token=HF_TOKEN)
            self._init_done = True

    async def initialize(self):
        """Initialize the vote service"""
        if self._initialized:
            await self._check_for_new_votes()
            return
        
        try:
            logger.info(LogFormatter.section("VOTE SERVICE INITIALIZATION"))
            
            # Ensure votes directory exists
            self.votes_file.parent.mkdir(parents=True, exist_ok=True)
            
            # Load existing votes if file exists
            local_vote_count = 0
            if self.votes_file.exists():
                logger.info(LogFormatter.info(f"Loading votes from {self.votes_file}"))
                local_vote_count = await self._count_local_votes()
                logger.info(LogFormatter.info(f"Found {local_vote_count:,} local votes"))
            
            # Check remote votes count
            remote_vote_count = await self._count_remote_votes()
            logger.info(LogFormatter.info(f"Found {remote_vote_count:,} remote votes"))
            
            if remote_vote_count > local_vote_count:
                logger.info(LogFormatter.info(f"Fetching {remote_vote_count - local_vote_count:,} new votes"))
                await self._sync_with_hub()
            elif remote_vote_count < local_vote_count:
                logger.warning(LogFormatter.warning(f"Local votes ({local_vote_count:,}) > Remote votes ({remote_vote_count:,})"))
                await self._load_existing_votes()
            else:
                logger.info(LogFormatter.success("Local and remote votes are in sync"))
                if local_vote_count > 0:
                    await self._load_existing_votes()
                else:
                    logger.info(LogFormatter.info("No votes found"))
            
            self._initialized = True
            self._last_sync = datetime.now(timezone.utc)
            
            # Final summary
            stats = {
                "Total_Votes": self._total_votes,
                "Last_Sync": self._last_sync.strftime("%Y-%m-%d %H:%M:%S UTC")
            }
            logger.info(LogFormatter.section("INITIALIZATION COMPLETE"))
            for line in LogFormatter.stats(stats):
                logger.info(line)
            
        except Exception as e:
            logger.error(LogFormatter.error("Initialization failed", e))
            raise

    async def _count_local_votes(self) -> int:
        """Count votes in local file"""
        if not self.votes_file.exists():
            return 0
            
        count = 0
        try:
            with open(self.votes_file, 'r') as f:
                for _ in f:
                    count += 1
            return count
        except Exception as e:
            logger.error(f"Error counting local votes: {str(e)}")
            return 0

    async def _count_remote_votes(self) -> int:
        """Count votes in remote file"""
        url = f"https://huggingface.co/datasets/{HF_ORGANIZATION}/votes/raw/main/votes_data.jsonl"
        headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
        
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(url, headers=headers) as response:
                    if response.status == 200:
                        count = 0
                        async for line in response.content:
                            if line.strip():  # Skip empty lines
                                count += 1
                        return count
                    else:
                        logger.error(f"Failed to get remote votes: HTTP {response.status}")
                        return 0
        except Exception as e:
            logger.error(f"Error counting remote votes: {str(e)}")
            return 0

    async def _sync_with_hub(self):
        """Sync votes with HuggingFace hub using datasets"""
        try:
            logger.info(LogFormatter.section("VOTE SYNC"))
            self._log_repo_operation("sync", f"{HF_ORGANIZATION}/votes", "Syncing local votes with HF hub")
            logger.info(LogFormatter.info("Syncing with HuggingFace hub..."))
            
            # Load votes from HF dataset
            dataset = datasets.load_dataset(
                f"{HF_ORGANIZATION}/votes", 
                split="train",
                cache_dir=cache_config.get_cache_path("datasets")
            )
            
            remote_votes = len(dataset)
            logger.info(LogFormatter.info(f"Dataset loaded with {remote_votes:,} votes"))
            
            # Convert to list of dictionaries
            df = dataset.to_pandas()
            if 'timestamp' in df.columns:
                df['timestamp'] = df['timestamp'].dt.strftime('%Y-%m-%dT%H:%M:%SZ')
            remote_votes = df.to_dict('records')
            
            # If we have more remote votes than local
            if len(remote_votes) > self._total_votes:
                new_votes = len(remote_votes) - self._total_votes
                logger.info(LogFormatter.info(f"Processing {new_votes:,} new votes..."))
                
                # Save votes to local file
                with open(self.votes_file, 'w') as f:
                    for vote in remote_votes:
                        f.write(json.dumps(vote) + '\n')
                
                # Reload votes in memory
                await self._load_existing_votes()
                logger.info(LogFormatter.success("Sync completed successfully"))
            else:
                logger.info(LogFormatter.success("Local votes are up to date"))

            self._last_sync = datetime.now(timezone.utc)
            
        except Exception as e:
            logger.error(LogFormatter.error("Sync failed", e))
            raise

    async def _check_for_new_votes(self):
        """Check for new votes on the hub"""
        try:
            self._log_repo_operation("check", f"{HF_ORGANIZATION}/votes", "Checking for new votes")
            # Load only dataset metadata
            dataset_info = datasets.load_dataset(f"{HF_ORGANIZATION}/votes", split="train")
            remote_vote_count = len(dataset_info)
            
            if remote_vote_count > self._total_votes:
                logger.info(f"Found {remote_vote_count - self._total_votes} new votes on hub")
                await self._sync_with_hub()
            else:
                logger.info("No new votes found on hub")
                
        except Exception as e:
            logger.error(f"Error checking for new votes: {str(e)}")

    async def _load_existing_votes(self):
        """Load existing votes from file"""
        if not self.votes_file.exists():
            logger.warning(LogFormatter.warning("No votes file found"))
            return

        try:
            logger.info(LogFormatter.section("LOADING VOTES"))
            
            # Clear existing data structures
            self.vote_check_set.clear()
            self._votes_by_model.clear()
            self._votes_by_user.clear()
            
            vote_count = 0
            latest_timestamp = None
            
            with open(self.votes_file, "r") as f:
                for line in f:
                    try:
                        vote = json.loads(line.strip())
                        vote_count += 1
                        
                        # Track latest timestamp
                        try:
                            vote_timestamp = datetime.fromisoformat(vote["timestamp"].replace("Z", "+00:00"))
                            if not latest_timestamp or vote_timestamp > latest_timestamp:
                                latest_timestamp = vote_timestamp
                            vote["timestamp"] = vote_timestamp.strftime("%Y-%m-%dT%H:%M:%SZ")
                        except (KeyError, ValueError) as e:
                            logger.warning(LogFormatter.warning(f"Invalid timestamp in vote: {str(e)}"))
                            continue
                        
                        if vote_count % 1000 == 0:
                            logger.info(LogFormatter.info(f"Processed {vote_count:,} votes..."))
                        
                        self._add_vote_to_memory(vote)
                        
                    except json.JSONDecodeError as e:
                        logger.error(LogFormatter.error("Vote parsing failed", e))
                        continue
                    except Exception as e:
                        logger.error(LogFormatter.error("Vote processing failed", e))
                        continue
            
            self._total_votes = vote_count
            self._last_vote_timestamp = latest_timestamp
            
            # Final summary
            stats = {
                "Total_Votes": vote_count,
                "Latest_Vote": latest_timestamp.strftime("%Y-%m-%d %H:%M:%S UTC") if latest_timestamp else "None",
                "Unique_Models": len(self._votes_by_model),
                "Unique_Users": len(self._votes_by_user)
            }
            
            logger.info(LogFormatter.section("VOTE SUMMARY"))
            for line in LogFormatter.stats(stats):
                logger.info(line)
            
        except Exception as e:
            logger.error(LogFormatter.error("Failed to load votes", e))
            raise

    def _add_vote_to_memory(self, vote: Dict[str, Any]):
        """Add vote to memory structures"""
        try:
            check_tuple = (vote["model"], vote["revision"], vote["username"])
            
            # Skip if we already have this vote
            if check_tuple in self.vote_check_set:
                return
                
            self.vote_check_set.add(check_tuple)
            
            # Update model votes
            if vote["model"] not in self._votes_by_model:
                self._votes_by_model[vote["model"]] = []
            self._votes_by_model[vote["model"]].append(vote)
            
            # Update user votes
            if vote["username"] not in self._votes_by_user:
                self._votes_by_user[vote["username"]] = []
            self._votes_by_user[vote["username"]].append(vote)
            
        except KeyError as e:
            logger.error(f"Malformed vote data, missing key: {str(e)}")
        except Exception as e:
            logger.error(f"Error adding vote to memory: {str(e)}")

    async def get_user_votes(self, user_id: str) -> List[Dict[str, Any]]:
        """Get all votes from a specific user"""
        logger.info(LogFormatter.info(f"Fetching votes for user: {user_id}"))
        votes = self._votes_by_user.get(user_id, [])
        logger.info(LogFormatter.success(f"Found {len(votes):,} votes"))
        return votes

    async def get_model_votes(self, model_id: str) -> Dict[str, Any]:
        """Get all votes for a specific model"""
        logger.info(LogFormatter.info(f"Fetching votes for model: {model_id}"))
        votes = self._votes_by_model.get(model_id, [])
        
        # Group votes by revision
        votes_by_revision = {}
        for vote in votes:
            revision = vote["revision"]
            if revision not in votes_by_revision:
                votes_by_revision[revision] = 0
            votes_by_revision[revision] += 1
        
        stats = {
            "Total_Votes": len(votes),
            **{f"Revision_{k}": v for k, v in votes_by_revision.items()}
        }
        
        logger.info(LogFormatter.section("VOTE STATISTICS"))
        for line in LogFormatter.stats(stats):
            logger.info(line)
        
        return {
            "total_votes": len(votes),
            "votes_by_revision": votes_by_revision,
            "votes": votes
        }

    async def _get_model_revision(self, model_id: str) -> str:
        """Get current revision of a model with retries"""
        logger.info(f"Getting revision for model: {model_id}")
        for attempt in range(self._max_retries):
            try:
                model_info = await asyncio.to_thread(self.hf_api.model_info, model_id)
                logger.info(f"Successfully got revision {model_info.sha} for model {model_id}")
                return model_info.sha
            except Exception as e:
                logger.error(f"Error getting model revision for {model_id} (attempt {attempt + 1}): {str(e)}")
                if attempt < self._max_retries - 1:
                    retry_delay = self._retry_delay * (attempt + 1)
                    logger.info(f"Retrying in {retry_delay} seconds...")
                    await asyncio.sleep(retry_delay)
                else:
                    logger.warning(f"Using 'main' as fallback revision for {model_id} after {self._max_retries} failed attempts")
                    return "main"

    async def add_vote(self, model_id: str, user_id: str, vote_type: str) -> Dict[str, Any]:
        """Add a vote for a model"""
        try:
            self._log_repo_operation("add", f"{HF_ORGANIZATION}/votes", f"Adding {vote_type} vote for {model_id} by {user_id}")
            logger.info(LogFormatter.section("NEW VOTE"))
            stats = {
                "Model": model_id,
                "User": user_id,
                "Type": vote_type
            }
            for line in LogFormatter.tree(stats, "Vote Details"):
                logger.info(line)
            
            revision = await self._get_model_revision(model_id)
            check_tuple = (model_id, revision, user_id)
            
            if check_tuple in self.vote_check_set:
                raise ValueError("Vote already recorded for this model")

            vote = {
                "model": model_id,
                "revision": revision,
                "username": user_id,
                "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
                "vote_type": vote_type
            }

            # Update local storage
            with open(self.votes_file, "a") as f:
                f.write(json.dumps(vote) + "\n")
            
            self._add_vote_to_memory(vote)
            self.votes_to_upload.append(vote)
            
            stats = {
                "Status": "Success",
                "Queue_Size": len(self.votes_to_upload)
            }
            for line in LogFormatter.stats(stats):
                logger.info(line)
            
            # Try to upload if batch size reached
            if len(self.votes_to_upload) >= self._upload_batch_size:
                logger.info(LogFormatter.info(f"Upload batch size reached ({self._upload_batch_size}), triggering sync"))
                await self._sync_with_hub()
            
            return {"status": "success", "message": "Vote added successfully"}
            
        except Exception as e:
            logger.error(LogFormatter.error("Failed to add vote", e))
            raise