#!/usr/bin/env python3
"""Extract all Docker image references from K8s YAML manifests."""
import os
import re
import sys
import fnmatch

REPO = "/home/ai-bot/Code/home-docker"
EXCLUDE_DIRS = ["archive", "telephone", "legacy-docker-compose", ".github", ".git"]

# Patterns to skip: :latest, no tag, or templated variables like ${...}
SKIP_TAG_PATTERNS = [
    re.compile(r'^latest$'),
    re.compile(r'^\$\{?.*\}?$'),
]

def should_exclude(path):
    rel = os.path.relpath(path, REPO)
    parts = rel.split(os.sep)
    for ex in EXCLUDE_DIRS:
        if ex in parts:
            return True
    return False

def extract_images(content):
    """Extract image references from YAML content (both image: field and image: line)."""
    images = set()
    for line in content.splitlines():
        # Match image: field in YAML
        m = re.match(r'^\s*image:\s*["\']?(.+?)["\']?\s*$', line)
        if not m:
            # Also match inline: image: "..." 
            m = re.search(r'image:\s*["\']?([a-zA-Z0-9._/-]+(?::[a-zA-Z0-9._-]+)?)["\']?', line)
            if m:
                img = m.group(1).strip()
                # Skip if it contains env var patterns
                if '${' in img or '$(' in img or img.startswith('$'):
                    continue
                images.add(img)
            continue
        img = m.group(1).strip().rstrip(',')
        # Handle YAML multiline/complex values
        img = img.strip('"').strip("'")
        # Skip if it contains env var patterns
        if '${' in img or '$(' in img or img.startswith('$'):
            continue
        # Skip empty
        if not img:
            continue
        images.add(img)
    return images

def parse_image(img_str):
    """Parse image string into (registy_prefix, image_name, tag)."""
    img_str = img_str.strip()
    # Check for explicit registry (docker.io, ghcr.io, quay.io, etc.)
    registry = ""
    
    if '/' in img_str:
        parts = img_str.split('/')
        # Check if first part looks like a registry (contains . or :port)
        if len(parts) >= 2 and ('.' in parts[0] or ':' in parts[0] or parts[0] == 'localhost' or parts[0] == 'docker.io'):
            registry = parts[0]
            image_path = '/'.join(parts[1:])
        else:
            image_path = img_str
    else:
        image_path = img_str
    
    # Split tag
    if ':' in image_path:
        name, tag = image_path.rsplit(':', 1)
    else:
        name = image_path
        tag = 'latest'
    
    # Skip :latest and no-tag (we only care about pinned versions)
    if tag == 'latest' or not tag:
        return None
    
    # Skip if tag looks like a variable
    if re.match(r'^\$\{', tag) or tag.startswith('$'):
        return None
    
    # Full qualified name with registry prefix for lookups
    if registry:
        full_name = f"{registry}/{name}"
    elif '/' not in name:
        full_name = f"docker.io/library/{name}"
    else:
        # Default to docker.io
        full_name = f"docker.io/{name}"
    
    return {
        'raw': img_str,
        'registry': registry or 'docker.io',
        'name': name,
        'tag': tag,
        'full_name': full_name,
        'lookup_name': f"{name}" if not registry else f"{registry}/{name}",
    }

def main():
    all_images = {}
    files_with_images = {}
    
    for root, dirs, files in os.walk(REPO):
        # Skip excluded dirs
        dirs[:] = [d for d in dirs if d not in EXCLUDE_DIRS and not d.startswith('.')]
        
        for fname in files:
            if not (fname.endswith('.yaml') or fname.endswith('.yml')):
                continue
            fpath = os.path.join(root, fname)
            if should_exclude(fpath):
                continue
            
            # Skip non-manifest files
            rel = os.path.relpath(fpath, REPO)
            if rel.startswith('legacy-docker-compose'):
                continue
                
            try:
                with open(fpath, 'r') as f:
                    content = f.read()
            except Exception as e:
                print(f"Error reading {fpath}: {e}", file=sys.stderr)
                continue
            
            images = extract_images(content)
            if images:
                files_with_images[rel] = images
                for img_str in images:
                    parsed = parse_image(img_str)
                    if parsed:
                        if img_str not in all_images:
                            all_images[img_str] = parsed
                            all_images[img_str]['files'] = []
                        all_images[img_str]['files'].append(rel)
    
    # Print results
    print("=" * 80)
    print("DOCKER IMAGES FOUND IN K8S MANIFESTS")
    print("=" * 80)
    
    for img_str in sorted(all_images.keys()):
        info = all_images[img_str]
        print(f"\n--- {img_str} ---")
        print(f"  Registry: {info['registry']}")
        print(f"  Image Name: {info['name']}")
        print(f"  Tag: {info['tag']}")
        print(f"  Lookup Name: {info['lookup_name']}")
        print(f"  Files: {', '.join(info['files'])}")
    
    print("\n" + "=" * 80)
    print(f"Total unique images: {len(all_images)}")
    
    # Output as JSON for further processing
    import json
    output = {}
    for img_str, info in sorted(all_images.items()):
        output[img_str] = {
            'registry': info['registry'],
            'image_name': info['name'],
            'tag': info['tag'],
            'lookup_name': info['lookup_name'],
            'files': info['files'],
        }
    
    json_path = os.path.join(REPO, 'scripts', 'images.json')
    with open(json_path, 'w') as f:
        json.dump(output, f, indent=2)
    print(f"\nJSON output written to: {json_path}")

if __name__ == '__main__':
    main()