#!/usr/bin/env python3

import argparse
import subprocess
import tempfile

import yaml


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Generate YAML for specified environment and image version."
    )
    parser.add_argument(
        "environment",
        choices=["prod", "staging"],
        help='Environment to use, either "prod" or "staging".',
    )
    return parser.parse_args()


def get_image_digest(image_name: str) -> str:
    get_digest_command = [
        "gcloud",
        "container",
        "images",
        "list-tags",
        image_name,
        "--filter=tags:latest",
        "--format=get(digest)",
    ]
    result = subprocess.run(get_digest_command, capture_output=True)
    # error handling (show stderr and error out)
    if result.returncode != 0:
        print(f"Error running {get_digest_command}. stderr:")
        print(result.stderr.decode())
        exit(1)
    digest = result.stdout.strip().decode()
    if len(digest) != 71:
        print(f"Could not extract digest from '{get_digest_command}' ")
        print("(should be 71 chars, beginning with sha256:). stdout:")
        print(result.stdout.decode())
        exit(1)
    return digest


def create_service_object(*, environment: str, digest: str) -> dict:
    hostname = "banking.withai.lol" if environment == "prod" else "staging.withai.lol"
    max_scale = 30 if environment == "prod" else 3
    min_scale = 1 if environment == "prod" else 0
    image_name = f"gcr.io/prompt-ad-game/promptgame-{environment}"
    throttling = "false" if environment == "prod" else "true"
    # Construct the underlying object as a nest of Python dicts/lists
    # Cloud Run service
    return {
        "apiVersion": "serving.knative.dev/v1",
        "kind": "Service",
        "metadata": {
            # Container name
            "name": f"promptgame-{environment}",
            # Our project name
            "namespace": "552241814655",
            "labels": {
                # Region
                "cloud.googleapis.com/location": "us-central1",
            },
        },
        "spec": {
            "template": {
                "metadata": {
                    "annotations": {
                        # Min/max instances
                        "autoscaling.knative.dev/minScale": str(min_scale),
                        "autoscaling.knative.dev/maxScale": str(max_scale),
                        # VPC egress settings
                        "run.googleapis.com/vpc-access-egress": "private-ranges-only",
                        "run.googleapis.com/vpc-access-connector": "serverless-vpc-promptgame",
                        # Allow ingress (incoming HTTP requests) from public internet
                        "run.googleapis.com/ingress": "all",
                        # Boost CPU during container spinup
                        "run.googleapis.com/startup-cpu-boost": "true",
                        # In prod, we always allocate a CPU and never spin down
                        # (throttling=false), while in staging we spin down to 0
                        # (throttling=true)
                        "run.googleapis.com/cpu-throttling": throttling,
                    },
                },
                "spec": {
                    # Max connections per container
                    "containerConcurrency": 300,
                    # Request timeout
                    "timeoutSeconds": 300,
                    # Service account name for container
                    "serviceAccountName": f"promptgame-service-worker@prompt-ad-game.iam.gserviceaccount.com",
                    "containers": [
                        # Image
                        {
                            "image": f"{image_name}@{digest}",
                            # Port to listen on
                            "ports": [{"name": "http1", "containerPort": 8080}],
                            "env": [
                                {
                                    "name": "APPLICATION_SETTINGS",
                                    "valueFrom": {
                                        "secretKeyRef": {
                                            "key": "latest",
                                            "name": f"promptgame_{environment}_application_settings",
                                        },
                                    },
                                }
                            ],
                            "resources": {
                                "limits": {
                                    "cpu": "1000m",
                                    "memory": "3G",
                                },
                            },
                            "startupProbe": {
                                "httpGet": {
                                    "path": "/robots.txt",
                                    "port": 8080,
                                    "httpHeaders": [
                                        {"name": "Host", "value": hostname}
                                    ],
                                },
                                "initialDelaySeconds": 10,
                                "timeoutSeconds": 10,
                                "failureThreshold": 10,
                                "periodSeconds": 10,
                            },
                        }
                    ],
                },
            },
            "traffic": [
                {
                    "percent": 100,
                    "latestRevision": True,
                }
            ],
        },
    }


def print_yaml(yaml_str: str) -> None:
    print("YAML config:")
    print("-" * 30)
    print(yaml_str)
    print("-" * 30)


def deploy_service(yaml_str: str) -> None:
    with tempfile.NamedTemporaryFile("w") as temp:
        temp.write(yaml_str)
        temp.flush()
        try:
            subprocess.run(
                ["gcloud", "run", "services", "replace", temp.name], check=True
            )
        except subprocess.CalledProcessError as e:
            print(f"Error running 'gcloud run services replace': {e}")
            exit(1)


def main() -> None:
    args = parse_arguments()
    image_name = f"gcr.io/prompt-ad-game/promptgame-{args.environment}"
    digest = get_image_digest(image_name)
    service_object = create_service_object(environment=args.environment, digest=digest)
    yaml_str = yaml.dump(service_object, default_flow_style=False)
    print_yaml(yaml_str)
    deploy_service(yaml_str)


if __name__ == "__main__":
    main()
