import json
import argparse
import subprocess
from pathlib import Path
from jinja2 import Template

### ========== GITOPS FUNCTIONS ==========

def run_command(command, cwd=None):
    result = subprocess.run(command, cwd=cwd, shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        raise Exception(f"Command '{command}' failed: {result.stderr}")
    return result.stdout.strip()

def clone(repo_url, clone_path: Path):
    run_command(f"git clone {repo_url} {clone_path}")

def checkout_base_branch(clone_path: Path, base_branch: str = "develop"):
    print(f"🔁 Checking out base branch: {base_branch}")
    run_command(f"git checkout {base_branch}", cwd=clone_path)

def checkout_branch(clone_path: Path, branch_name: str):
    run_command(f"git checkout -b {branch_name}", cwd=clone_path)

def add(clone_path: Path, file_path: Path):
    rel_path = file_path.relative_to(clone_path)
    run_command(f"git add {rel_path}", cwd=clone_path)

def commit(clone_path: Path, commit_message: str):
    run_command(f'git commit -m "{commit_message}"', cwd=clone_path)

def push(clone_path: Path, branch_name: str):
    run_command(f"git push origin {branch_name}", cwd=clone_path)

def create_file_from_template(file_path: Path, template_content: str, context: dict, dryrun: bool):
    template = Template(template_content)
    rendered = template.render(context)
    if dryrun:
        print(f"\n📝 Would write to {file_path}:\n{rendered}")
    else:
        file_path.parent.mkdir(parents=True, exist_ok=True)
        file_path.write_text(rendered)

def cleanup(clone_path: Path):
    run_command(f"rm -rf {clone_path}")

### ========== NODE OPERATION ==========

def process_node(vmname: str, ipaddress: str, gateway: str, clone_path: Path,
                 commit_template: str, file_template: str, dryrun: bool):

    file_rel_path = Path(f"hieradata/nodes/{vmname}.yaml")
    file_path = clone_path / file_rel_path
    branch_name = f"autonode/{vmname}"

    if file_path.exists() and not dryrun:
        print(f"⚠️  Skipping {vmname}: {file_path} already exists.")
        return

    print(f"\n🌿 Creating branch: {branch_name}")
    checkout_branch(clone_path, branch_name)

    print(f"📝 Rendering YAML for {vmname}")
    create_file_from_template(
        file_path,
        file_template,
        {
            "ipaddress": ipaddress,
            "gateway": gateway,
            "interface": "eth0"
        },
        dryrun
    )

    if dryrun:
        print(f"💤 Dry run: skipping add/commit/push for {vmname}")
        return

    print(f"➕ Adding {file_rel_path}")
    add(clone_path, file_path)

    commit_msg = Template(commit_template).render({"vmname": vmname})
    print(f"✅ Committing: {commit_msg}")
    commit(clone_path, commit_msg)

    print(f"🚀 Pushing {branch_name}")
    push(clone_path, branch_name)


def load_broken_tf_outputs(file_path: Path):
    """Handles newline-separated JSON objects (non-standard tf_outputs.json format)."""
    objects = []
    buffer = ""
    for line in file_path.read_text().splitlines():
        line = line.strip()
        if not line:
            continue
        buffer += line
        if buffer.endswith("}"):
            try:
                obj = json.loads(buffer)
                objects.append(obj)
                buffer = ""
            except json.JSONDecodeError:
                buffer += " "  # accumulate more lines until it's valid
    return objects

### ========== MAIN CLI SCRIPT ==========

def main():
    parser = argparse.ArgumentParser(description="Generate Hiera node YAMLs and push to Git")
    parser.add_argument("--output-json", required=True, type=Path, help="Terragrunt JSON outputs")
    parser.add_argument("--repo-url", required=True, help="Git repo URL")
    parser.add_argument("--clone-path", required=True, type=Path, help="Temp clone path")
    parser.add_argument("--commit-template", required=True, help="Commit message Jinja2 template")
    parser.add_argument("--file-template", required=True, type=Path, help="Path to Jinja2 YAML template")
    parser.add_argument("--dry-run", action="store_true", help="Do not write or push, just preview")
    parser.add_argument("--base-branch", default="develop", help="Base branch to branch off (default: develop)")

    args = parser.parse_args()

    if args.clone_path.exists():
        print(f"🧹 Removing existing clone at {args.clone_path}")
        cleanup(args.clone_path)

    print(f"📥 Cloning repo to {args.clone_path}")
    clone(args.repo_url, args.clone_path)

    file_template = args.file_template.read_text()
    # Use loader
    parsed_objects = load_broken_tf_outputs(args.output_json)

    # Flatten into merged format using hostnames
    merged_outputs = {}
    for obj in parsed_objects:
        if "vm_metadata" in obj and "value" in obj["vm_metadata"]:
            hostname = obj["vm_metadata"]["value"]["hostname"]
            merged_outputs[f"vm_metadata_{hostname}"] = obj["vm_metadata"]

    for module_path, data in merged_outputs.items():
        if "value" not in data:
            print(f"⏭️  Skipping {module_path}: missing 'value'")
            continue

        node = data["value"]
        vmname = node["hostname"]
        ip = node["ipaddress"]
        gw = node["gateway"]

        checkout_base_branch(args.clone_path, args.base_branch)

        print(f"\n🔧 Processing {vmname} ({ip})")
        process_node(
            vmname=vmname,
            ipaddress=ip,
            gateway=gw,
            clone_path=args.clone_path,
            commit_template=args.commit_template,
            file_template=file_template,
            dryrun=args.dry_run
        )

    if not args.dry_run:
        print(f"\n🧹 Cleaning up: {args.clone_path}")
        cleanup(args.clone_path)

    print("\n🏁 All done!")

if __name__ == "__main__":
    main()