from kubernetes import client, config


def update_max_allowed_vpa(api_instance, namespace, vpa_name, max_allowed):
    try:
        vpa = api_instance.get_namespaced_custom_object(
            "autoscaling.k8s.io",
            "v1",
            namespace,
            "verticalpodautoscalers",
            vpa_name,
        )
        # skip if updateMode is off
        if vpa["spec"]["updatePolicy"]["updateMode"] == "Off":
            print(f"  Skipping VPA '{vpa_name}' in namespace '{namespace}' because updateMode is Off")
            return
        vpa["spec"]["resourcePolicy"]["containerPolicies"][0]["maxAllowed"] = max_allowed
        api_instance.replace_namespaced_custom_object(
            "autoscaling.k8s.io",
            "v1",
            namespace,
            "verticalpodautoscalers",
            vpa_name,
            vpa
        )
        print(f"  Updated maxAllowed for VPA '{vpa_name}' in namespace '{namespace}'")
    except client.rest.ApiException as e:
        print(f"  Error updating VPA '{vpa_name}' in namespace '{namespace}': {e}")


def main():
    # Load Kubernetes configuration
    config.load_kube_config()

    # Define cluster names
    clusters = [
        # "nonprod.np.navi-tech.in",
        # "aps1.prod.navi-tech.in",
        "aps1.np.navi-gi.in",
        # "aps1.prod.navi-gi.in",
        # "aps1.np.navi-sa.in",
        # "aps1.prod.navi-sa.in",
        # "aps1.np.navi-pay.in",
        # "aps1.prod.navi-pay.in",
    ]

    max_allowed = {
        "cpu": "7200m",
        "memory": "16Gi"
    }

    for cluster in clusters:
        print(f"Processing cluster: {cluster}")
        try:
            api_instance = client.CustomObjectsApi()
            config.load_kube_config(context=cluster)
            v1 = client.CoreV1Api()
            ns = [ns.metadata.name for ns in v1.list_namespace().items]
            for namespace in ns:
                print(f"  Processing namespace: {namespace}")
                try:
                    vpas = api_instance.list_namespaced_custom_object(
                        "autoscaling.k8s.io",
                        "v1",
                        namespace,
                        "verticalpodautoscalers",
                    )
                    for vpa in vpas.get("items", []):
                        vpa_name = vpa.get("metadata", {}).get("name")
                        if vpa_name:
                            update_max_allowed_vpa(api_instance, namespace, vpa_name,
                                                   max_allowed)
                except client.rest.ApiException as e:
                    print(f"    Error listing VPA in namespace '{namespace}': {e}")
        except client.rest.ApiException as e:
            print(f"  Error processing cluster '{cluster}': {e}")


if __name__ == "__main__":
    main()