#!/usr/bin/env python3
"""
Check Ollama library for local models (without 'cloud' tag)
and print their names with size tags
"""
import requests
from bs4 import BeautifulSoup
import logging

logger = logging.getLogger(__name__)

# ANSI color codes
GREEN = "\033[32m"
RESET = "\033[0m"

OLLAMA_LIBRARY_URL = "https://ollama.com/library"


def fetch_ollama_models():
    """
    Fetch and parse Ollama library for local models
    
    Returns:
        list: List of model dictionaries with name and sizes (oldest first, newest last)
    """
    try:
        # Fetch with newest sort - this returns models sorted newest first
        url = f"{OLLAMA_LIBRARY_URL}?sort=newest"
        logger.info(f"Fetching Ollama library from {url}...")
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        
        soup = BeautifulSoup(response.text, 'html.parser')
        
        models = []
        
        # Find all model list items
        model_items = soup.find_all('li', attrs={'x-test-model': True})
        
        for item in model_items:
            # Get model name
            name_span = item.find('span', class_='group-hover:underline truncate')
            if not name_span:
                continue
            
            model_name = name_span.get_text(strip=True)

            # Get description
            desc_p = item.find('p', class_='max-w-lg break-words text-neutral-800 text-md')
            description = desc_p.get_text(strip=True).lower() if desc_p else ""

            # Check if description contains instruct-related terms
            is_instruct = 'instruct' in description

            # Find all tags in this model item using CSS selector
            tag_spans = item.select('span[class*="rounded-md"]')

            # Check if this model has a 'cloud' tag
            has_cloud = False
            sizes = []

            for tag in tag_spans:
                tag_text = tag.get_text(strip=True).lower()

                # Check for cloud tag (cyan colored)
                if tag_text == 'cloud':
                    has_cloud = True
                    break

                # Check for size tags (have x-test-size attribute)
                if tag.has_attr('x-test-size'):
                    sizes.append(tag_text)

            # Skip models with cloud tag
            if has_cloud:
                continue

            models.append({
                'name': model_name,
                'sizes': sizes,
                'is_instruct': is_instruct
            })
        
        # Reverse so oldest first, newest at the end
        return models[::-1]
        
    except Exception as e:
        logger.error(f"Error fetching Ollama data: {e}")
        return []


def print_models(models):
    """
    Print models in a readable format
    
    Args:
        models: List of model dictionaries
    """
    if not models:
        print("No local Ollama models found.")
        return

    total = len(models)
    print(f"\nFound {total} local Ollama models (without 'cloud' tag):\n")
    print("-" * 80)

    for i, model in enumerate(models, 1):
        name = model['name']
        sizes = model['sizes']
        is_instruct = model.get('is_instruct', False)
        # Number from newest (1) to oldest (total)
        num = total - i + 1

        # Build tags string
        tags = []
        if is_instruct:
            tags.append(f"{GREEN}INSTRUCT{RESET}")
        if sizes:
            tags.extend(sizes)

        if tags:
            tags_str = ', '.join(tags)
            print(f"{num}. {name} ({tags_str})")
        else:
            print(f"{num}. {name}")

    print("-" * 80)
    print(f"\nNewest models at top (lower numbers), oldest at bottom (higher numbers)")


def main():
    """
    Main function to fetch and display local Ollama models
    """
    logging.basicConfig(level=logging.WARNING)
    
    models = fetch_ollama_models()
    print_models(models)


if __name__ == "__main__":
    main()
