#!/usr/bin/env python3
import argparse
import json
import subprocess
import sys
import textwrap


def get_project():
    """Get the currently selected project."""
    result = subprocess.run(
        ["gcloud", "config", "get-value", "project"],
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        text=True,
        check=True,
    )
    return result.stdout.strip()


def run_gcloud(command_args):
    """Runs a gcloud command with specified arguments and returns parsed JSON output."""
    full_command = ["gcloud"] + command_args + ["--format=json"]
    print(f"Running command: {' '.join(full_command)}")
    result = subprocess.run(full_command, capture_output=True, text=True, check=True)
    return json.loads(result.stdout)


def get_project_resources(project):
    """Get all needed resources from the project."""
    project_args = ["--project", project]
    vpcs = run_gcloud(["compute", "networks", "list"] + project_args)
    subnets = run_gcloud(["compute", "networks", "subnets", "list"] + project_args)
    firewall_rules = run_gcloud(["compute", "firewall-rules", "list"] + project_args)
    instances = run_gcloud(["compute", "instances", "list"] + project_args)
    config = run_gcloud(["compute", "project-info", "describe"])
    project_resources = {
        "vpcs": vpcs,
        "subnets": subnets,
        "firewall_rules": firewall_rules,
        "instances": instances,
        "config": config,
    }

    return project_resources


def find_resource(project_resources, self_link):
    """Find a resource by self-link"""
    for resource_type, resources in project_resources.items():
        for resource in resources:
            if resource["selfLink"] == self_link:
                return resource
    raise ValueError("Resource not found")


def main():
    parser = argparse.ArgumentParser(description="GCP Project Resource Checkup Tool")
    parser.add_argument(
        "--save", help="Store the project resources dictionary to a JSON file"
    )
    parser.add_argument(
        "--load", help="Load project resources from a JSON file and skip gcloud calls"
    )
    args = parser.parse_args()

    if args.load:
        try:
            with open(args.load, "r") as f:
                project_resources = json.load(f)
            print(f"--- Loaded data from {args.load} ---")
        except Exception as e:
            print(f"Error loading file {args.load}: {e}")
            sys.exit(1)
    else:
        project = get_project()
        if not project:
            print(
                textwrap.dedent("""
            You don't have a project selected. Please run the command:
            
              gcloud config set project <your-project-id>

            NOTE: Your project ID is NOT cis-91. It's something like cis-91-123823
            """)
            )
            sys.exit(1)
        project_resources = get_project_resources(project)

    if args.save:
        with open(args.save, "w") as f:
            json.dump(project_resources, f, indent=2)
        print(f"--- Data saved to {args.save} ---")
        sys.exit(0)

    vpc = None

    print("Checking your project metadata...")
    metadata = {
        item["key"]: item["value"]
        for item in project_resources["config"]["commonInstanceMetadata"]["items"]
    }
    assert "ssh-keys" in metadata, (
        """You don't have any SSH keys loaded into the Compute Engine metadata."""
    )
    assert "startup-script" in metadata, (
        """You don't have a startup script in Compute Engine. The startup script should install and configure your website."""
    )

    print("Checking VMs:")
    for instance in project_resources["instances"]:
        print("-->", instance["name"])
        assert len(instance["serviceAccounts"]) > 0, (
            """There are no service accounts attached to your VM"""
        )
        print("  + Service account:", instance["serviceAccounts"][0]["email"])
        assert "-compute" not in instance["serviceAccounts"][0]["email"], (
            """Your VM is using the default service account!"""
        )
        vpc = find_resource(
            project_resources, instance["networkInterfaces"][0]["network"]
        )
        print("  + VPC:", vpc["name"])
        assert vpc["name"] != "default", (
            """Your VM is on the default VPC. It should be on a custom VPC."""
        )
        print("  + Tags:")
        for tag in instance["tags"]["items"]:
            print("    - ", tag)
        assert len(instance["tags"]["items"]) > 0, (
            """Your VM should have at least one tag (to match the firewall configuration.)"""
        )
        assert len(instance["networkInterfaces"]) > 0, (
            """Your VM doesn't have a network interface"""
        )
        assert "natIP" in instance["networkInterfaces"][0]["accessConfigs"][0], (
            """Your VM doesn't have a public IP"""
        )
        public_ip = instance["networkInterfaces"][0]["accessConfigs"][0]["natIP"]
        print("  + Public IP:", public_ip)
        print("  + Checking the web page on the VM")
        web = subprocess.run(
            f"curl -m 2 http://{public_ip}",
            encoding="utf-8",
            capture_output=True,
            shell=True,
        )
        assert 0 == web.returncode, f"""Failed to fetch http://{public_ip}"""
        assert "You are visitor" in web.stdout, (
            """The web page didn't report the visit count."""
        )

    assert vpc is not None, """You don't have any VMs!"""

    print(f"Checking subnets in VPC {vpc['name']}:")
    subnets = [
        net for net in project_resources["subnets"] if net["network"] == vpc["selfLink"]
    ]
    regions = [
        net["region"].split("/")[-1]
        for net in project_resources["subnets"]
        if net["network"] == vpc["selfLink"]
    ]
    assert 3 == len(subnets), """You should have exactly three subnets."""
    for subnet in subnets:
        print(
            "-->",
            subnet["name"],
            ":",
            subnet["ipCidrRange"],
            "in region",
            subnet["region"].split("/")[-1],
        )
    assert set(regions) == set(["us-central1", "us-west1", "us-east1"]), (
        """You should have a subnet in the us-central1, us-west1, and us-east1 regions."""
    )

    print(f"Firewall rules in VPC {vpc['name']}:")
    firewall_rules = [
        rule
        for rule in project_resources["firewall_rules"]
        if rule["network"] == vpc["selfLink"]
    ]
    for rule in firewall_rules:
        print("-->", rule["name"])
        if "allowed" in rule:
            for allow in rule["allowed"]:
                ports = allow.get("ports", ["ALL"])
                print("  + protocol:", allow["IPProtocol"], "ports:", ",".join(ports))

        print("  +", "direction:", rule["direction"])
        print("  +", "source:", ", ".join(rule["sourceRanges"]))
        if "targetTags" in rule:
            print("  +", "target tags:", ", ".join(rule["targetTags"]))
        else:
            print("  +", "target: ALL")


if __name__ == "__main__":
    main()