import json
import os
import requests
from copy import deepcopy
from pathlib import Path
import pytest
import time

api_key = os.environ.get("API_KEY")
port = os.environ.get("PORT", 9001)
base_url = os.environ.get("BASE_URL", "http://localhost")

with open(os.path.join(Path(__file__).resolve().parent, "sam_tests.json"), "r") as f:
    TESTS = json.load(f)


@pytest.mark.skip(reason="SAM testing is broken, to be fixed in future release")
@pytest.mark.parametrize("test", TESTS)
def test_sam(test):
    payload = deepcopy(test["payload"])
    payload["api_key"] = api_key
    response = requests.post(
        f"{base_url}:{port}/sam/{test['type']}",
        json=payload,
    )
    try:
        response.raise_for_status()
        data = response.json()
        if test["type"] == "embed_image":
            try:
                assert "embeddings" in data
            except:
                print(f"Invalid response: {data}, expected 'embeddings' in response")
            try:
                assert len(data["embeddings"]) == len(
                    test["expected_response"]["embeddings"]
                )
            except:
                print(
                    f"Invalid response: {data}, expected length of embeddings to be {len(test['expected_response']['embeddings'])}, got {len(data['embeddings'])}"
                )
        if test["type"] == "segment_image":
            try:
                assert "masks" in data
            except:
                print(f"Invalid response: {data}, expected 'masks' in response")
            try:
                assert data["masks"] == test["expected_response"]["masks"]
            except:
                print(
                    f"Invalid response: {data}, expected masks to be {test['expected_response']['masks']}, got {data['masks']}"
                )
    except Exception as e:
        raise e


@pytest.fixture(scope="session", autouse=True)
def setup():
    try:
        res = requests.get(f"{base_url}:{port}")
        res.raise_for_status()
        success = True
    except:
        success = False

    waited = 0
    while not success:
        print("Waiting for server to start...")
        time.sleep(5)
        waited += 5
        try:
            res = requests.get(f"{base_url}:{port}")
            res.raise_for_status()
            success = True
        except:
            success = False
        if waited > 30:
            raise Exception("Test server failed to start")


if __name__ == "__main__":
    test_sam()