import os
import openai
import unittest

from openai.error import AuthenticationError, RateLimitError

# No API key is provided implicitly, to avoid leakage
assert os.environ.get('OPENAI_API_KEY') is None

# This is my API key
MINE = os.environ['MY_OPENAI_API_KEY']

prompt = "Give me three anagrams of 'Randompile':"


class NoKeyError(Exception):
    pass


def inspector(f):
    def _(*a, **kw):
        try:
            result = f(*a, **kw)
        except RateLimitError:
            raise RuntimeError("This test assumes the API key is still valid.")
        except AuthenticationError as ex:
            if "No API key provided" in ex.user_message:
                raise NoKeyError
            else:
                raise
        else:
            return result

    return _


@inspector
def completion(*, key_constr=None, key_member=None):
    worker = openai.Completion(api_key=key_constr)
    return worker.create(api_key=key_member, model='text-davinci-003', prompt=prompt)


@inspector
def dalleimage(*, key_constr=None, key_member=None):
    worker = openai.Image(api_key=key_constr)
    return worker.create(api_key=key_member, prompt=prompt)


class ObservedBehavior(unittest.TestCase):
    def test_completion_00(self):
        with self.assertRaises(NoKeyError):
            completion(key_constr=None, key_member=None)

    def test_completion_10_UNEXPECTED_BEHAVIOUR(self):
        with self.assertRaises(NoKeyError):
            completion(key_constr=MINE, key_member=None)

    def test_completion_01(self):
        result = completion(key_constr=None, key_member=MINE)
        # print(f"{prompt}", result['choices'][0]['text'].strip(), sep='\n')

    def test_completion_11(self):
        result = completion(key_constr=MINE, key_member=MINE)
        # print(f"{prompt}", result['choices'][0]['text'].strip(), sep='\n')

    def test_dalleimage_00(self):
        with self.assertRaises(NoKeyError):
            dalleimage(key_constr=None, key_member=None)

    def test_dalleimage_10_UNEXPECTED_BEHAVIOUR(self):
        with self.assertRaises(NoKeyError):
            dalleimage(key_constr=MINE, key_member=None)

    def test_dalleimage_01_UNEXPECTED_BEHAVIOUR(self):
        with self.assertRaises(NoKeyError):
            dalleimage(key_constr=None, key_member=MINE)

    def test_dalleimage_11_UNEXPECTED_BEHAVIOUR(self):
        with self.assertRaises(NoKeyError):
            dalleimage(key_constr=MINE, key_member=MINE)


if __name__ == '__main__':
    unittest.main()