#!/usr/bin/env python3
import requests
import json
from collections import Counter

def get_model_counts():
    # Base URL of the API endpoint
    url = "https://civitai.com/api/v1/models"
    
    # Parameters for the API request
    params = {
        'limit': 100,  # Set a high limit to reduce the number of pages to iterate through
        'page': 1      # Start from the first page
    }

    # Initialize counts
    total_models = 0

    # counter
    allow_commercial_use = []
    creator = []

    # sum
    nsfw_count = []
    allow_derivatives = []
    allow_no_credit = []
    download_count = []
    favorite_count = []
    comment_count = []
    rating_count = []
    tipped_amount_count = []
    ratings = []
    failures = 0
    total_pages = 1000 # set to very high value, just in case it fails a lot in the beginning

    while True:
        # Make the GET request
        try:
            response = requests.get(url, params=params)
            response.raise_for_status()  # This will raise an error if the request fails
            data = response.json()

            # Process the current page of results
            models = data["items"]
            total_models += len(models)

            for model in models:
                allow_commercial_use.append(model["allowCommercialUse"])
                creator.append(model["creator"]["username"])
                nsfw_count.append(model["nsfw"])
                allow_derivatives.append(model["allowDerivatives"])
                allow_no_credit.append(model["allowNoCredit"])
                download_count.append(model["stats"]["downloadCount"])
                favorite_count.append(model["stats"]["favoriteCount"])
                comment_count.append(model["stats"]["commentCount"])
                rating_count.append(model["stats"]["ratingCount"])
                tipped_amount_count.append(model["stats"]["tippedAmountCount"])
                ratings.append(model["stats"]["rating"])

            # Check if there are more pages to process
            total_pages = int(data.get('metadata', {}).get('totalPages', 0))
            # if params['page'] >= total_pages:
            if params['page'] >= 2:
                break

            # Increment the page number for the next request
            params['page'] += 1

            print(f"{params['page']} / {total_pages}")
        except:
            failures += 1
            params["page"] += 1

    return {
        'total_models': total_models,
        'allow_commercial_use': allow_commercial_use,
        'creator': creator,
        'nsfw_count': nsfw_count,
        'allow_derivatives': allow_derivatives,
        'allow_no_credit': allow_no_credit,
        'download_count': download_count,
        'favorite_count': favorite_count,
        'comment_count': comment_count,
        'rating_count': rating_count,
        'tipped_amount_count': tipped_amount_count,
        'ratings': ratings,
        'failures': failures,
    }

outputs = get_model_counts()

def map_fn(k, v):
    if k in ["total_models", 'failures']:
        return v
    elif k in ["creator", "allow_commercial_use"]:
        return Counter(v)
    elif k in ["ratings"]:
        return sum(v) / len(v)
    else:
        return {
            "sum": sum(v),
            "avg": round(sum(v) / len(v), 3)
        }

stats = {k: map_fn(k, v) for k,v in outputs.items()}
stats["num_creators"] = len(stats.pop("creator"))

for k in ["total_models", "download_count", "comment_count", "rating_count", "tipped_amount_count", "favorite_count"]:
    if isinstance(stats[k], dict):
        stats[f"{k}_per_creator"] = round(stats[k]["sum"] / stats["num_creators"], 2)
    else:
        stats[f"{k}_per_creator"] = round(stats[k] / stats["num_creators"], 2)

filename = "civitai_stats.json"
with open(filename, 'w') as file:
    json.dump(stats, file, indent=4)