#!/usr/bin/env python3
"""
Generate Ansible playbooks from existing docker-compose files in the homelab repo.
This script scans the hosts/ directory and creates deployment playbooks.
"""

import os
import yaml
from pathlib import Path
from collections import defaultdict

REPO_ROOT = Path(__file__).parent.parent.parent
HOSTS_DIR = REPO_ROOT / "hosts"
ANSIBLE_DIR = Path(__file__).parent
PLAYBOOKS_DIR = ANSIBLE_DIR / "playbooks"
HOST_VARS_DIR = ANSIBLE_DIR / "host_vars"

# Mapping of directory names to ansible host names
HOST_MAPPING = {
    "atlantis": "atlantis",
    "calypso": "calypso",
    "setillo": "setillo",
    "guava": "guava",
    "concord-nuc": "concord_nuc",
    "anubis": "anubis",
    "homelab-vm": "homelab_vm",
    "chicago-vm": "chicago_vm",
    "bulgaria-vm": "bulgaria_vm",
    "contabo-vm": "contabo_vm",
    "rpi5-vish": "rpi5_vish",
    "tdarr-node": "tdarr_node",
}

# Host categories for grouping
HOST_CATEGORIES = {
    "synology": ["atlantis", "calypso", "setillo"],
    "physical": ["guava", "concord-nuc", "anubis"],
    "vms": ["homelab-vm", "chicago-vm", "bulgaria-vm", "contabo-vm", "matrix-ubuntu-vm"],
    "edge": ["rpi5-vish", "nvidia_shield"],
    "proxmox": ["tdarr-node"],
}


def find_compose_files():
    """Find all docker-compose files in the hosts directory."""
    compose_files = defaultdict(list)

    for yaml_file in HOSTS_DIR.rglob("*.yaml"):
        if ".git" in str(yaml_file):
            continue
        compose_files[yaml_file.parent].append(yaml_file)

    for yml_file in HOSTS_DIR.rglob("*.yml"):
        if ".git" in str(yml_file):
            continue
        compose_files[yml_file.parent].append(yml_file)

    return compose_files


def get_host_from_path(file_path):
    """Extract REDACTED_APP_PASSWORD path."""
    parts = file_path.relative_to(HOSTS_DIR).parts

    # Structure: hosts/<category>/<host>/...
    if len(parts) >= 2:
        category = parts[0]
        host = parts[1]
        return category, host
    return None, None


def extract_service_name(file_path):
    """Extract service name from file path."""
    # Get the service name from parent directory or filename
    if file_path.name in ["docker-compose.yml", "docker-compose.yaml"]:
        return file_path.parent.name
    else:
        return file_path.stem.replace("-", "_").replace(".", "_")


def is_compose_file(file_path):
    """Check if file looks like a docker-compose file."""
    try:
        with open(file_path, 'r') as f:
            content = yaml.safe_load(f)
            if content and isinstance(content, dict):
                return 'services' in content or 'version' in content
    except:
        pass
    return False


def generate_service_vars(host, services):
    """Generate host_vars with service definitions."""
    service_list = []

    for service_path, service_name in services:
        rel_path = service_path.relative_to(REPO_ROOT)

        # Determine the stack directory name
        if service_path.name in ["docker-compose.yml", "docker-compose.yaml"]:
            stack_dir = service_path.parent.name
        else:
            stack_dir = service_name

        service_entry = {
            "name": service_name,
            "stack_dir": stack_dir,
            "compose_file": str(rel_path),
            "enabled": True,
        }

        # Check for .env file
        env_file = service_path.parent / ".env"
        stack_env = service_path.parent / "stack.env"
        if env_file.exists():
            service_entry["env_file"] = str(env_file.relative_to(REPO_ROOT))
        elif stack_env.exists():
            service_entry["env_file"] = str(stack_env.relative_to(REPO_ROOT))

        service_list.append(service_entry)

    return service_list


def generate_host_playbook(host_name, ansible_host, services, category):
    """Generate a playbook for a specific host."""

    # Create header comment
    header = f"""---
# Deployment playbook for {host_name}
# Category: {category}
# Services: {len(services)}
#
# Usage:
#   ansible-playbook playbooks/deploy_{ansible_host}.yml
#   ansible-playbook playbooks/deploy_{ansible_host}.yml -e "stack_deploy=false"
#   ansible-playbook playbooks/deploy_{ansible_host}.yml --check

"""

    playbook = [
        {
            "name": f"Deploy services to {host_name}",
            "hosts": ansible_host,
            "gather_facts": True,
            "vars": {
                "services": "{{ host_services | default([]) }}"
            },
            "tasks": [
                {
                    "name": "Display deployment info",
                    "ansible.builtin.debug": {
                        "msg": "Deploying {{ services | length }} services to {{ inventory_hostname }}"
                    }
                },
                {
                    "name": "Ensure docker data directory exists",
                    "ansible.builtin.file": {
                        "path": "{{ docker_data_path }}",
                        "state": "directory",
                        "mode": "0755"
                    }
                },
                {
                    "name": "Deploy each enabled service",
                    "ansible.builtin.include_role": {
                        "name": "docker_stack"
                    },
                    "vars": {
                        "stack_name": "{{ item.stack_dir }}",
                        "stack_compose_file": "{{ item.compose_file }}",
                        "stack_env_file": "{{ item.env_file | default(omit) }}"
                    },
                    "loop": "{{ services }}",
                    "loop_control": {
                        "label": "{{ item.name }}"
                    },
                    "when": "item.enabled | default(true)"
                }
            ]
        }
    ]

    return header, playbook


def main():
    """Main function to generate all playbooks."""
    print("=" * 60)
    print("Generating Ansible Playbooks from Homelab Repository")
    print("=" * 60)

    # Ensure directories exist
    PLAYBOOKS_DIR.mkdir(parents=True, exist_ok=True)
    HOST_VARS_DIR.mkdir(parents=True, exist_ok=True)

    # Find all compose files
    compose_files = find_compose_files()

    # Organize by host
    hosts_services = defaultdict(list)

    for directory, files in compose_files.items():
        category, host = get_host_from_path(directory)
        if not host:
            continue

        for f in files:
            if is_compose_file(f):
                service_name = extract_service_name(f)
                hosts_services[(category, host)].append((f, service_name))

    # Generate playbooks and host_vars
    all_hosts = {}

    for (category, host), services in sorted(hosts_services.items()):
        ansible_host = HOST_MAPPING.get(host, host.replace("-", "_"))

        print(f"\n[{category}/{host}] Found {len(services)} services:")
        for service_path, service_name in services:
            print(f"  - {service_name}")

        # Generate host_vars
        service_vars = generate_service_vars(host, services)
        host_vars = {
            "host_services": service_vars
        }

        host_vars_file = HOST_VARS_DIR / f"{ansible_host}.yml"
        with open(host_vars_file, 'w') as f:
            f.write("---\n")
            f.write(f"# Auto-generated host variables for {host}\n")
            f.write(f"# Services deployed to this host\n\n")
            yaml.dump(host_vars, f, default_flow_style=False, sort_keys=False)

        # Generate individual host playbook
        header, playbook = generate_host_playbook(host, ansible_host, services, category)
        playbook_file = PLAYBOOKS_DIR / f"deploy_{ansible_host}.yml"
        with open(playbook_file, 'w') as f:
            f.write(header)
            yaml.dump(playbook, f, default_flow_style=False, sort_keys=False)

        all_hosts[ansible_host] = {
            "category": category,
            "host": host,
            "services": len(services)
        }

    # Generate master playbook
    master_playbook = [
        {
            "name": "Deploy all homelab services",
            "hosts": "localhost",
            "gather_facts": False,
            "tasks": [
                {
                    "name": "Display deployment plan",
                    "ansible.builtin.debug": {
                        "msg": "Deploying services to all hosts. Use --limit to target specific hosts."
                    }
                }
            ]
        }
    ]

    # Add imports for each host
    for ansible_host, info in sorted(all_hosts.items()):
        master_playbook.append({
            "name": f"Deploy to {info['host']} ({info['services']} services)",
            "ansible.builtin.import_playbook": f"playbooks/deploy_{ansible_host}.yml",
            "tags": [info['category'], ansible_host]
        })

    master_file = ANSIBLE_DIR / "site.yml"
    with open(master_file, 'w') as f:
        f.write("---\n")
        f.write("# Master Homelab Deployment Playbook\n")
        f.write("# Auto-generated from docker-compose files\n")
        f.write("#\n")
        f.write("# Usage:\n")
        f.write("#   Deploy everything:     ansible-playbook site.yml\n")
        f.write("#   Deploy specific host:  ansible-playbook site.yml --limit atlantis\n")
        f.write("#   Deploy by category:    ansible-playbook site.yml --tags synology\n")
        f.write("#\n\n")
        yaml.dump(master_playbook, f, default_flow_style=False, sort_keys=False)

    print(f"\n{'=' * 60}")
    print(f"Generated playbooks for {len(all_hosts)} hosts")
    print(f"Master playbook: {master_file}")
    print("=" * 60)


if __name__ == "__main__":
    main()
