import json
import os
import requests
import subprocess
import yaml

DOCS_PATH = '/docs' if os.environ.get('CALICO_DOCS_PATH') is None else os.environ.get('CALICO_DOCS_PATH')
RELEASE_STREAM = os.environ.get('RELEASE_STREAM')
EXCLUDED_IMAGES = ['calico/pilot-webhook',
                   'calico/upgrade',
                   'quay.io/coreos/flannel']

GCR_IMAGES = ['calico/node',
              'calico/cni',
              'calico/typha']
EXPECTED_ARCHS = ['amd64', 'arm64', 'ppc64le']

VERSIONS_WITHOUT_FLANNEL_MIGRATION = ['v3.8', 'v3.7', 'v3.6', 'v3.5', 'v3.4', 'v3.3', 'v3.2', 'v3.1', 'v3.0']
if RELEASE_STREAM in VERSIONS_WITHOUT_FLANNEL_MIGRATION:
    EXCLUDED_IMAGES.append('calico/flannel-migration-controller')
    print('[INFO] excluding "calico/flannel-migration-controller" for older release')

with open('%s/_data/versions.yml' % DOCS_PATH) as f:
    versions = yaml.safe_load(f)
    RELEASE_VERSION = versions[0]['title']
    print('[INFO] using _data/versions.yaml, discovered version: %s' % RELEASE_VERSION)

def test_operator_image_present():
    with open('%s/_data/versions.yml' % DOCS_PATH) as versionsFile:
        versions = yaml.safe_load(versionsFile)
        for version in versions:
            if version["title"] == RELEASE_VERSION:
                # Found matching version. Perform the test.
                operator = version["tigera-operator"]
                img = "%s/%s:%s" % (operator["registry"], operator["image"], operator["version"])
                print('[INFO] checking %s' % img)
                headers = {'content-type': 'application/json'}
                req = requests.get("https://quay.io/api/v1/repository/tigera/operator/tag/%s/images" % (
                    operator["version"]), headers=headers)
                assert req.status_code == 200
                return
        assert False, "Unable to find matching version"

def test_quay_release_tag_present():
    with open('%s/_config.yml' % DOCS_PATH) as config:
        images = yaml.safe_load(config)
        for image in images['imageNames']:
            if images['imageNames'][image] not in EXCLUDED_IMAGES:
                print('[INFO] checking quay.io/%s:%s' % (images['imageNames'][image], RELEASE_VERSION))

                headers = {'content-type': 'application/json'}
                req = requests.get("https://quay.io/api/v1/repository/%s/tag/%s/images" % (
                    images['imageNames'][image], RELEASE_VERSION), headers=headers)
                assert req.status_code == 200


def test_gcr_release_tag_present():
    with open('%s/_config.yml' % DOCS_PATH) as config:
        images = yaml.safe_load(config)
        for image in images['imageNames']:
            if images['imageNames'][image] in GCR_IMAGES:
                gcr_name = images['imageNames'][image].replace('calico/', '')
                print('[INFO] checking gcr.io/projectcalico-org/%s:%s' % (gcr_name, RELEASE_VERSION))
                cmd = 'docker manifest inspect gcr.io/projectcalico-org/%s:%s | jq -r "."' % (gcr_name, RELEASE_VERSION)

                req = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
                try:
                    metadata = json.loads(req.stdout.read())
                except ValueError:
                    print("[ERROR] Didn't get json back from docker manifest inspect.  Does image exist?")
                    assert False
                found_archs = []
                for platform in metadata['manifests']:
                    found_archs.append(platform['platform']['architecture'])

                assert EXPECTED_ARCHS == found_archs


def test_docker_release_tag_present():
    with open('%s/_config.yml' % DOCS_PATH) as config:
        images = yaml.safe_load(config)
        for image in images['imageNames']:
            if images['imageNames'][image] not in EXCLUDED_IMAGES:
                print('[INFO] checking %s:%s' % (images['imageNames'][image], RELEASE_VERSION))
                cmd = 'docker manifest inspect %s:%s | jq -r "."' % (images['imageNames'][image], RELEASE_VERSION)

                req = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
                metadata = json.loads(req.stdout.read())
                found_archs = []
                for platform in metadata['manifests']:
                    found_archs.append(platform['platform']['architecture'])

                assert EXPECTED_ARCHS.sort() == found_archs.sort()

@@ -2048,33 +2048,43 @@
 %5D:%0A            i
-f
+mage_name =
  images%5B'imageNa
@@ -2087,32 +2087,84 @@
 geNames'%5D%5Bimage%5D
+.replace(%22docker.io/%22, %22%22)%0A            if image_name
  not in EXCLUDED
@@ -2234,38 +2234,21 @@
 %25 (image
-s%5B'imageNames'%5D%5Bimage%5D
+_name
 , RELEAS
@@ -2440,38 +2440,21 @@
    image
-s%5B'imageNames'%5D%5Bimage%5D
+_name
 , RELEAS
@@ -2709,33 +2709,43 @@
 %5D:%0A            i
-f
+mage_name =
  images%5B'imageNa
@@ -2756,16 +2756,68 @@
 %5D%5Bimage%5D
+.replace(%22docker.io/%22, %22%22)%0A            if image_name
  in GCR_
@@ -2856,38 +2856,21 @@
  = image
-s%5B'imageNames'%5D%5Bimage%5D
+_name
 .replace
@@ -3853,25 +3853,35 @@
            i
-f
+mage_name =
  images%5B'ima
@@ -3896,16 +3896,68 @@
 %5D%5Bimage%5D
+.replace(%22docker.io/%22, %22%22)%0A            if image_name
  not in 
@@ -4027,38 +4027,21 @@
 %25 (image
-s%5B'imageNames'%5D%5Bimage%5D
+_name
 , RELEAS
@@ -4126,38 +4126,21 @@
 %25 (image
-s%5B'imageNames'%5D%5Bimage%5D
+_name
 , RELEAS

Fix image tests now that names are prefixed with docker.io