import json
import subprocess
import os
import hashlib

def get_hash(file_path):
    if not os.path.exists(file_path):
        return "MISSING"
    sha256_hash = hashlib.sha256()
    with open(file_path,"rb") as f:
        for byte_block in iter(lambda: f.read(4096),b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()

def main():
    # 1. Get ollama list
    try:
        raw_list = subprocess.check_output(["ollama", "list"]).decode("utf-8")
    except:
        raw_list = ""

    lines = raw_list.strip().split("\n")[1:]
    inventory = []

    for line in lines:
        parts = line.split()
        if len(parts) < 3: continue
        tag = parts[0]
        tag_id = parts[1]
        size = parts[2]

        # 2. Find manifest path
        # Pattern: /var/lib/ollama/.ollama/models/manifests/registry.ollama.ai/library/NAME/TAG
        # Or: /var/lib/ollama/.ollama/models/manifests/registry.ollama.ai/USER/NAME/TAG

        manifest_root = "/var/lib/ollama/.ollama/models/manifests/registry.ollama.ai/"
        tag_parts = tag.split("/")
        if len(tag_parts) == 1:
            # library
            name_tag = tag_parts[0].split(":")
            name = name_tag[0]
            version = name_tag[1] if len(name_tag) > 1 else "latest"
            manifest_path = os.path.join(manifest_root, "library", name, version)
        else:
            user = tag_parts[0]
            name_tag = tag_parts[1].split(":")
            name = name_tag[0]
            version = name_tag[1] if len(name_tag) > 1 else "latest"
            manifest_path = os.path.join(manifest_root, user, name, version)

        manifest_hash = get_hash(manifest_path)

        # 3. Read manifest content
        layers = []
        config_digest = ""
        if os.path.exists(manifest_path):
            try:
                with open(manifest_path, "r") as f:
                    data = json.load(f)
                    config_digest = data.get("config", {}).get("digest", "")
                    for layer in data.get("layers", []):
                        layers.append({
                            "mediaType": layer.get("mediaType"),
                            "digest": layer.get("digest"),
                            "size": layer.get("size"),
                            "from": layer.get("from")
                        })
            except:
                pass

        # 4. Get Modelfile for critical params
        try:
            modelfile = subprocess.check_output(["ollama", "show", tag, "--modelfile"]).decode("utf-8")
            from_line = [l for l in modelfile.split("\n") if l.startswith("FROM ")][0]
        except:
            modelfile = ""
            from_line = ""

        entry = {
            "tag": tag,
            "ollama_list_id": tag_id,
            "ollama_list_size": size,
            "manifest_path": manifest_path,
            "manifest_hash": manifest_hash,
            "config_digest": config_digest,
            "blob_digest_from": from_line.replace("FROM ", "").strip(),
            "layers": layers,
            "reconstructible": "YES" if from_line and manifest_hash != "MISSING" else "UNKNOWN"
        }
        inventory.append(entry)

    output = {
        "inventory_date": "2026-05-25T13:35:00",
        "total_tags": len(inventory),
        "models": inventory
    }

    print(json.dumps(output, indent=2))

if __name__ == "__main__":
    main()