#!/usr/bin/env python3
"""
Utility to parse ollama ls output and extract model names.

Replicates the Jinja2 filter logic from:
  ollama_existing_models.stdout_lines
  | reject('match', '^\\s*$|^(NAME|\\-+)')
  | map('regex_replace', '^\\s*(\\S+)\\s.*', '\\1')
  | list

Usage:
    python parse_ollama_models.py < input.txt
    python parse_ollama_models.py models_output.txt
    echo "model_name ..." | python parse_ollama_models.py
"""

import re
import sys


def parse_ollama_models(lines):
    """Parse ollama ls output lines and return list of model names."""
    results = []
    for line in lines:
        # Skip empty lines and header rows (NAME, ---)
        if not line.strip() or re.match(r'^(NAME|\-+)$', line.strip()):
            continue
        # Extract the model name (first non-whitespace token)
        match = re.match(r'^\s*(\S+)\s.*', line)
        if match:
            results.append(match.group(1))
    return results


def main():
    # Read input from file argument, stdin, or skip
    if len(sys.argv) > 1:
        with open(sys.argv[1], 'r') as f:
            lines = f.read().splitlines()
    else:
        lines = sys.stdin.read().splitlines()

    models = parse_ollama_models(lines)

    # Output: one model per line
    for model in models:
        print(model)

    # Also write to a JSON file if a second arg is given
    if len(sys.argv) > 2:
        import json
        with open(sys.argv[2], 'w') as f:
            json.dump(models, f, indent=2)


if __name__ == '__main__':
    main()