import os
import re
import time
import spaces
import requests
import gradio as gr
from huggingface_hub import HfApi, list_models
from threading import Thread, Lock

HF_TOKEN = os.environ.get("HF_TOKEN", None)

TITLE = """
<h1><center>Open-Schizo-Leaderboard</center></h1>
<center>
<p>Comparing LLM Cards for how absolutely Schizo they are</p>
<p>If you like my work please subscribe for $5 a month to <a href="https://www.patreon.com/Rombodawg" target="_blank">https://www.patreon.com/Rombodawg</a></p>
</center>
"""

CSS = """
.duplicate-button {
    margin: auto !important;
    color: white !important;
    background: black !important;
    border-radius: 100vh !important;
}
h3 {
    text-align: center;
}
table {
    width: 100%;
    border-collapse: collapse;
}
table, th, td {
    border: 1px solid #ddd;
}
th, td {
    padding: 8px;
    text-align: left;
}
th {
    background-color: #f2f2f2;
    cursor: pointer;
}
tr:nth-child(even) {
    background-color: #f9f9f9;
}
tr:hover {
    background-color: #f1f1f1;
}
.leaderboard-container {
    max-height: 600px;
    overflow-y: auto;
}
.loading {
    text-align: center;
    font-size: 18px;
    padding: 20px;
}
"""

SCHIZO_WORDS = [
    "MAXED", "Max", "SUPER", "Duped", "Edge", "maid", "Solution", 
    "gpt-4", "gpt4o", "claude-3.5", "claude-3.7", "o1", "o3-mini", 
    "gpt-4.5", "chatgpt", "merge", "merged", "best", "greatest", 
    "highest quality", "Class 1", "NSFW", "4chan", "reddit", "vibe", 
    "vibe check", "vibe checking", "dirty", "meme", "memes", "upvote", 
    "Linear", "SLERP", "Nearswap", "Task Arithmetic", "Task_Arithmetic", 
    "TIES", "DARE", "Passthrough", "Model Breadcrumbs", "Model Stock", 
    "NuSLERP", "DELL", "DELLA Task Arithmeti", "SCE"
]

MARKDOWN_SYMBOLS = ["#", "*", "_", "`", ">", "-", "+", "[", "]", "(", ")", "!", "\\", "|", "~", "<", ">", "=", ":"]

CACHE = {
    "llm": {"html": None, "timestamp": 0, "lock": Lock()},
    "all": {"html": None, "timestamp": 0, "lock": Lock()}
}

def count_schizo_words(text):
    count = 0
    for word in SCHIZO_WORDS:
        count += len(re.findall(re.escape(word), text, re.IGNORECASE))
    return count

def count_markdown_symbols(text):
    count = 0
    for symbol in MARKDOWN_SYMBOLS:
        count += text.count(symbol)
    return count

def calculate_word_count(text):
    return len(re.findall(r'\w+', text))

def calculate_schizo_rating(readme_content):
    schizo_word_count = count_schizo_words(readme_content)
    word_schizo_rating = schizo_word_count * 10
    
    word_count = calculate_word_count(readme_content)
    wordiness_schizo_rating = 0
    if word_count < 150:
        wordiness_schizo_rating = word_schizo_rating * 0.5
    elif word_count > 1000:
        extra_penalty = min(1.0, 0.5 + ((word_count - 1000) // 500) * 0.25)
        wordiness_schizo_rating = word_schizo_rating * extra_penalty
    
    markdown_count = count_markdown_symbols(readme_content)
    visual_schizo_rating = 0
    if markdown_count > 100:
        visual_penalty = min(1.0, 0.5 + ((markdown_count - 150) // 50) * 0.25)
        visual_schizo_rating = word_schizo_rating * visual_penalty
    
    combined_schizo_rating = word_schizo_rating + wordiness_schizo_rating + visual_schizo_rating
    
    return {
        "combined": combined_schizo_rating,
        "word": word_schizo_rating,
        "wordiness": wordiness_schizo_rating,
        "visual": visual_schizo_rating,
        "schizo_word_count": schizo_word_count,
        "word_count": word_count,
        "markdown_count": markdown_count
    }

def fetch_model_readme(model_id):
    try:
        url = f"https://huggingface.co/{model_id}/raw/main/README.md"
        response = requests.get(url)
        return response.text if response.status_code == 200 else None
    except Exception as e:
        print(f"Error fetching README for {model_id}: {e}")
        return None

def generate_leaderboard_data(model_type="llm", max_models=500):
    api = HfApi(token=HF_TOKEN)
    
    models = list_models(
        task="text-generation" if model_type == "llm" else None,
        limit=max_models
    )
    
    leaderboard_data = []
    for model in models:
        readme_content = fetch_model_readme(model.id)
        if not readme_content:
            continue
        
        ratings = calculate_schizo_rating(readme_content)
        leaderboard_data.append({
            "model_id": model.id,
            "combined_rating": ratings["combined"],
            "word_rating": ratings["word"],
            "wordiness_rating": ratings["wordiness"],
            "visual_rating": ratings["visual"],
            "schizo_word_count": ratings["schizo_word_count"],
            "word_count": ratings["word_count"],
            "markdown_count": ratings["markdown_count"]
        })
    
    leaderboard_data.sort(key=lambda x: x["combined_rating"], reverse=True)
    return leaderboard_data

def create_leaderboard_html(leaderboard_data):
    html = """
    <div class="leaderboard-container">
        <table id="leaderboard">
            <tr>
                <th onclick="sortTable(0)">Model</th>
                <th onclick="sortTable(1, true)">Average Schizo Rating</th>
                <th onclick="sortTable(2, true)">Visual Schizo Rating</th>
                <th onclick="sortTable(3, true)">Wordiness Schizo Rating</th>
                <th onclick="sortTable(4, true)">Overall Schizo Rating</th>
            </tr>
    """
    
    for item in leaderboard_data:
        html += f"""
            <tr>
                <td>{item["model_id"]}</td>
                <td>{item["combined_rating"]:.2f}</td>
                <td>{item["visual_rating"]:.2f}</td>
                <td>{item["wordiness_rating"]:.2f}</td>
                <td>{item["word_rating"]:.2f}</td>
            </tr>
        """
    
    html += """
        </table>
    </div>
    <script>
    function sortTable(n, isNumeric = false) {
        var table, rows, switching, i, x, y, shouldSwitch, dir, switchcount = 0;
        table = document.getElementById("leaderboard");
        switching = true;
        dir = "asc";
        
        while (switching) {
            switching = false;
            rows = table.rows;
            
            for (i = 1; i < (rows.length - 1); i++) {
                shouldSwitch = false;
                x = rows[i].getElementsByTagName("TD")[n];
                y = rows[i + 1].getElementsByTagName("TD")[n];
                
                if (dir == "asc") {
                    if (isNumeric) {
                        if (parseFloat(x.innerHTML) > parseFloat(y.innerHTML)) {
                            shouldSwitch = true;
                            break;
                        }
                    } else {
                        if (x.innerHTML.toLowerCase() > y.innerHTML.toLowerCase()) {
                            shouldSwitch = true;
                            break;
                        }
                    }
                } else if (dir == "desc") {
                    if (isNumeric) {
                        if (parseFloat(x.innerHTML) < parseFloat(y.innerHTML)) {
                            shouldSwitch = true;
                            break;
                        }
                    } else {
                        if (x.innerHTML.toLowerCase() < y.innerHTML.toLowerCase()) {
                            shouldSwitch = true;
                            break;
                        }
                    }
                }
            }
            
            if (shouldSwitch) {
                rows[i].parentNode.insertBefore(rows[i + 1], rows[i]);
                switching = true;
                switchcount++;
            } else {
                if (switchcount == 0 && dir == "asc") {
                    dir = "desc";
                    switching = true;
                }
            }
        }
    }
    </script>
    """
    return html

def background_cache_updater():
    while True:
        try:
            for model_type in ["llm", "all"]:
                with CACHE[model_type]["lock"]:
                    leaderboard_data = generate_leaderboard_data(model_type)
                    CACHE[model_type]["html"] = create_leaderboard_html(leaderboard_data)
                    CACHE[model_type]["timestamp"] = time.time()
        except Exception as e:
            print(f"Background update error: {e}")
        time.sleep(600)

def load_leaderboard(model_type):
    with CACHE[model_type]["lock"]:
        if CACHE[model_type]["html"]:
            return CACHE[model_type]["html"]
        return f'<div class="loading">{model_type.upper()} leaderboard loading...</div>'

def trigger_refresh(model_type):
    def refresh_task():
        try:
            with CACHE[model_type]["lock"]:
                leaderboard_data = generate_leaderboard_data(model_type)
                CACHE[model_type]["html"] = create_leaderboard_html(leaderboard_data)
                CACHE[model_type]["timestamp"] = time.time()
        except Exception as e:
            print(f"Refresh error: {e}")
    Thread(target=refresh_task, daemon=True).start()
    return f'<div class="loading">{model_type.upper()} refresh started...</div>'

@spaces.GPU()
def init_leaderboard():
    return '<div class="loading">Initializing leaderboard...</div>'

with gr.Blocks(css=CSS, theme="soft") as demo:
    gr.HTML(TITLE)
    gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
    
    with gr.Row():
        with gr.Column():
            model_type_dropdown = gr.Dropdown(
                choices=["llm", "all"],
                value="llm",
                label="Model Type Filter",
            )
            refresh_button = gr.Button("Force Refresh Current View")
        
    leaderboard_html = gr.HTML(init_leaderboard())
    
    model_type_dropdown.change(
        fn=load_leaderboard,
        inputs=[model_type_dropdown],
        outputs=[leaderboard_html]
    )
    
    demo.load(
        fn=load_leaderboard,
        inputs=[model_type_dropdown],
        outputs=[leaderboard_html],
        every=30
    )
    
    refresh_button.click(
        fn=trigger_refresh,
        inputs=[model_type_dropdown],
        outputs=[leaderboard_html],
    )

if __name__ == "__main__":
    Thread(target=background_cache_updater, daemon=True).start()
    demo.launch()