model_suggestions.py

from nc_py_api import Nextcloud
import json
from typing import Dict, List
import os
import time
from datetime import datetime
import threading
import arena_config
import re

# Initialize Nextcloud client
nc = Nextcloud(
    nextcloud_url=arena_config.NEXTCLOUD_URL,
    nc_auth_user=arena_config.NEXTCLOUD_USERNAME,
    nc_auth_pass=arena_config.NEXTCLOUD_PASSWORD
)

SUGGESTIONS_FILE = "model_suggestions.json"
NEXTCLOUD_SUGGESTIONS_PATH = "/gpu_poor_model_suggestions.json"

def validate_model_url(url: str) -> bool:
    """Validate if the provided URL matches the expected HuggingFace format."""
    pattern = r'^hf\.co/[\w-]+/[\w\.-]+(?:-GGUF)?:Q[0-9]+(?:_[A-Z0-9_]+)?$'
    return bool(re.match(pattern, url))

def load_suggestions() -> Dict:
    """Load suggestions from Nextcloud, fallback to local file, or initialize if neither exists."""
    suggestions = None
    
    # First try to load from Nextcloud
    try:
        remote_data = nc.files.download(NEXTCLOUD_SUGGESTIONS_PATH)
        if remote_data:
            suggestions = json.loads(remote_data.decode('utf-8'))
            # Update local file with Nextcloud data
            with open(SUGGESTIONS_FILE, 'w') as f:
                json.dump(suggestions, f, indent=2)
            return suggestions
    except Exception as e:
        print(f"Could not load from Nextcloud: {e}")
    
    # If Nextcloud fails, try local file
    if os.path.exists(SUGGESTIONS_FILE):
        try:
            with open(SUGGESTIONS_FILE, 'r') as f:
                suggestions = json.load(f)
                return suggestions
        except Exception as e:
            print(f"Could not load from local file: {e}")
    
    # If both fail, initialize new suggestions
    suggestions = {
        "suggestions": {},
        "last_updated": datetime.now().isoformat()
    }
    
    # Save the new suggestions to both local and Nextcloud
    save_suggestions(suggestions)
    
    return suggestions

def save_suggestions(suggestions: Dict) -> None:
    """Save suggestions to both local file and Nextcloud."""
    with open(SUGGESTIONS_FILE, 'w') as f:
        json.dump(suggestions, f, indent=2)
    
    # Upload to Nextcloud
    try:
        nc.files.upload(
            NEXTCLOUD_SUGGESTIONS_PATH,
            json.dumps(suggestions, indent=2).encode('utf-8')
        )
    except Exception as e:
        print(f"Error uploading to Nextcloud: {e}")

def add_suggestion(model_url: str) -> str:
    """Add a new model suggestion or update existing one."""
    if not validate_model_url(model_url):
        return "❌ Invalid model URL format. Please use the format: hf.co/username/model-name-GGUF:Q4_K_M"
    
    suggestions = load_suggestions()
    
    if model_url in suggestions["suggestions"]:
        suggestions["suggestions"][model_url]["count"] += 1
        suggestions["suggestions"][model_url]["last_suggested"] = datetime.now().isoformat()
        message = f"✨ Model suggestion updated! This model has been suggested {suggestions['suggestions'][model_url]['count']} times."
    else:
        suggestions["suggestions"][model_url] = {
            "count": 1,
            "first_suggested": datetime.now().isoformat(),
            "last_suggested": datetime.now().isoformat()
        }
        message = "✅ New model suggestion recorded successfully!"
    
    suggestions["last_updated"] = datetime.now().isoformat()
    save_suggestions(suggestions)
    
    return message

def get_suggestions_html() -> str:
    """Generate HTML table of model suggestions."""
    suggestions = load_suggestions()
    
    # Sort suggestions by count (descending)
    sorted_suggestions = sorted(
        suggestions["suggestions"].items(),
        key=lambda x: x[1]["count"],
        reverse=True
    )
    
    html = """
    <style>
        .suggestions-table {
            width: 100%;
            border-collapse: collapse;
            font-family: Arial, sans-serif;
        }
        .suggestions-table th, .suggestions-table td {
            border: 1px solid #ddd;
            padding: 8px;
            text-align: left;
        }
        .suggestions-table th {
            background-color: rgba(255, 255, 255, 0.1);
            font-weight: bold;
        }
        .rank-column {
            width: 60px;
            text-align: center;
        }
        .count-badge {
            background-color: rgba(34, 87, 122, 0.7);
            color: white;
            padding: 4px 8px;
            border-radius: 12px;
            font-size: 0.9em;
        }
        .description-column {
            font-size: 0.9em;
            color: #888;
        }
    </style>
    <table class='suggestions-table'>
    <tr>
        <th class='rank-column'>Rank</th>
        <th>Model URL</th>
        <th>Suggestions</th>
        <th>First Suggested</th>
        <th>Last Suggested</th>
    </tr>
    """
    
    for index, (model_url, data) in enumerate(sorted_suggestions, start=1):
        rank_display = {1: "🥇", 2: "🥈", 3: "🥉"}.get(index, f"{index}")
        
        html += f"""
        <tr>
            <td class='rank-column'>{rank_display}</td>
            <td>{model_url}</td>
            <td><span class="count-badge">{data['count']}</span></td>
            <td>{data['first_suggested'].split('T')[0]}</td>
            <td>{data['last_suggested'].split('T')[0]}</td>
        </tr>
        """
    
    html += "</table>"
    return html