import subprocess
import json
import os
import glob

def get_models_from_modelfiles():
    models = set()
    modelfiles = glob.glob("Modelfile.*")
    for modelfile in modelfiles:
        with open(modelfile, 'r') as f:
            for line in f:
                if line.strip().startswith("FROM"):
                    parts = line.strip().split()
                    if len(parts) > 1:
                        model_name = parts[1]
                        # Only consider models that are not local file paths
                        if not model_name.startswith('/') and not model_name.startswith('~') and not model_name.endswith('.gguf'):
                            models.add(model_name)
                    break # Only take the first FROM in each Modelfile
    return sorted(list(models))

def update_ollama_model(model_name):
    print(f"Checking for updates for model: {model_name}")
    try:
        # Run ollama pull command
        process = subprocess.run(
            ["ollama", "pull", model_name],
            capture_output=True,
            text=True,
            check=True,
            timeout=900 # 15 minutes
        )
        output = process.stdout
        print(f"Output for {model_name}:\n{output}")

        # Basic check to see if an update happened.
        # Ollama pull output will contain "pulling" or "downloading" if an update is in progress
        # and "success" if it completed. If the model is already up to date, it says "already up to date".
        if "pulling" in output or "downloading" in output:
            print(f"Model {model_name} was updated.")
            return True
        elif "already up to date" in output:
            print(f"Model {model_name} is already up to date.")
            return False
        else:
            print(f"Unexpected output for {model_name}, assuming no update: {output}")
            return False

    except subprocess.CalledProcessError as e:
        print(f"Error updating model {model_name}: {e}")
        print(f"Stderr: {e.stderr}")
        return False
    except FileNotFoundError:
        print("Error: 'ollama' command not found. Please ensure Ollama is installed and in your PATH.")
        return False

def main():
    models_to_update = get_models_from_modelfiles()
    print(f"Identified models to check for updates: {models_to_update}")

    updated_models = []
    for model in models_to_update:
        if update_ollama_model(model):
            updated_models.append(model)

    if updated_models:
        print("\nSuccessfully updated the following models:")
        for model in updated_models:
            print(f"- {model}")
    else:
        print("\nNo models were updated.")

if __name__ == "__main__":
    main()