#!/usr/bin/env python3
"""Generate Ansible playbooks from plain English descriptions using Ollama LLM.

Usage:
  python3 scripts/ansible-generator.py "Install and configure nginx on all debian hosts"
  python3 scripts/ansible-generator.py --name nginx-setup "Install nginx on debian hosts"
  python3 scripts/ansible-generator.py --dry-run "Restart Docker on atlantis and calypso"
"""

import argparse
import logging
import re
import sys
from pathlib import Path

import yaml

sys.path.insert(0, str(Path(__file__).parent))
from lib.ollama import ollama_generate, ollama_available, OllamaUnavailableError

log = logging.getLogger(__name__)

REPO_DIR = Path("/home/homelab/organized/repos/homelab")
PLAYBOOKS_DIR = REPO_DIR / "ansible" / "playbooks"
GENERATED_DIR = PLAYBOOKS_DIR / "generated"
INVENTORY_PATH = REPO_DIR / "ansible" / "inventory.yml"

# Reference playbooks for style
STYLE_REFS = [
    PLAYBOOKS_DIR / "deploy_atlantis.yml",
    PLAYBOOKS_DIR / "update_system.yml",
]


def load_style_examples() -> str:
    """Load first 50 lines of reference playbooks for style guidance."""
    examples = []
    for ref in STYLE_REFS:
        if ref.exists():
            lines = ref.read_text().splitlines()[:50]
            examples.append(f"# {ref.name}\n" + "\n".join(lines))
    return "\n\n".join(examples)


def load_inventory() -> str:
    if INVENTORY_PATH.exists():
        return INVENTORY_PATH.read_text()
    return "(inventory not found)"


def slugify(text: str) -> str:
    """Convert description to a filename-safe slug."""
    slug = re.sub(r"[^a-z0-9]+", "-", text.lower().strip())
    slug = slug.strip("-")
    return slug[:60] if slug else "playbook"


def generate_playbook(description: str) -> str:
    examples = load_style_examples()
    inventory = load_inventory()

    prompt = (
        f"Generate an Ansible playbook for: {description}\n\n"
        f"Follow this style:\n{examples}\n\n"
        f"Use this inventory:\n{inventory}\n\n"
        "Output valid YAML only. Include proper hosts, tasks, handlers, and tags.\n"
        "Do not wrap the YAML in code fences. Output only the raw YAML document starting with ---."
    )
    return ollama_generate(prompt, num_predict=3000, timeout=180)


def extract_yaml(raw: str) -> str:
    """Strip any markdown code fences the LLM may have added."""
    # Remove ```yaml ... ``` wrapping
    raw = re.sub(r"^```(?:ya?ml)?\s*\n", "", raw.strip())
    raw = re.sub(r"\n```\s*$", "", raw.strip())
    return raw.strip()


def validate_yaml(content: str) -> bool:
    """Check that content is valid YAML."""
    try:
        yaml.safe_load(content)
        return True
    except yaml.YAMLError as e:
        log.error("Generated YAML is invalid: %s", e)
        return False


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate Ansible playbooks from plain English via LLM")
    parser.add_argument("description", help="Plain English description of the desired playbook")
    parser.add_argument("--name", type=str, help="Output filename (without .yml extension)")
    parser.add_argument("--dry-run", action="store_true", help="Print playbook without saving")
    parser.add_argument("--verbose", "-v", action="store_true", help="Enable debug logging")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    if not ollama_available():
        log.error("Ollama is not reachable — aborting")
        sys.exit(1)

    log.info("Generating playbook for: %s", args.description)
    raw = generate_playbook(args.description)
    playbook_yaml = extract_yaml(raw)

    if not validate_yaml(playbook_yaml):
        log.error("LLM output failed YAML validation. Raw output:\n%s", raw)
        sys.exit(1)

    if args.dry_run:
        print(playbook_yaml)
        return

    filename = args.name if args.name else slugify(args.description)
    if not filename.endswith(".yml"):
        filename += ".yml"

    GENERATED_DIR.mkdir(parents=True, exist_ok=True)
    out_path = GENERATED_DIR / filename
    out_path.write_text(playbook_yaml + "\n")
    log.info("Saved playbook to %s", out_path)


if __name__ == "__main__":
    main()