import os
import openai
openai.api_type = "azure"
openai.api_base = os.getenv("OPENAI_API_BASE")
openai.api_version = "2023-03-15-preview"
openai.api_key = os.getenv("OPENAI_API_KEY")

messages = []
response = None
pronpt = None
code = None
code_string = None

openapi_spec = '''{
  "swagger": "2.0",
  "info": {
    "version": "1.0.0",
    "title": "Swagger Petstore",
    "description": "A sample API that uses a petstore as an example to demonstrate features in the swagger-2.0 specification",
    "termsOfService": "http://swagger.io/terms/",
    "contact": {
      "name": "Swagger API Team"
    },
    "license": {
      "name": "MIT"
    }
  },
  "host": "petstore.swagger.io",
  "basePath": "/api",
  "schemes": [
    "http"
  ],
  "consumes": [
    "application/json"
  ],
  "produces": [
    "application/json"
  ],
  "paths": {
    "/pets": {
      "get": {
        "description": "Returns all pets from the system that the user has access to",
        "operationId": "findPets",
        "produces": [
          "application/json",
          "application/xml",
          "text/xml",
          "text/html"
        ],
        "parameters": [
          {
            "name": "tags",
            "in": "query",
            "description": "tags to filter by",
            "required": false,
            "type": "array",
            "items": {
              "type": "string"
            },
            "collectionFormat": "csv"
          },
          {
            "name": "limit",
            "in": "query",
            "description": "maximum number of results to return",
            "required": false,
            "type": "integer",
            "format": "int32"
          }
        ],
        "responses": {
          "200": {
            "description": "pet response",
            "schema": {
              "type": "array",
              "items": {
                "$ref": "#/definitions/Pet"
              }
            }
          },
          "default": {
            "description": "unexpected error",
            "schema": {
              "$ref": "#/definitions/ErrorModel"
            }
          }
        }
      },
      "post": {
        "description": "Creates a new pet in the store.  Duplicates are allowed",
        "operationId": "addPet",
        "produces": [
          "application/json"
        ],
        "parameters": [
          {
            "name": "pet",
            "in": "body",
            "description": "Pet to add to the store",
            "required": true,
            "schema": {
              "$ref": "#/definitions/NewPet"
            }
          }
        ],
        "responses": {
          "200": {
            "description": "pet response",
            "schema": {
              "$ref": "#/definitions/Pet"
            }
          },
          "default": {
            "description": "unexpected error",
            "schema": {
              "$ref": "#/definitions/ErrorModel"
            }
          }
        }
      }
    },
    "/pets/{id}": {
      "get": {
        "description": "Returns a user based on a single ID, if the user does not have access to the pet",
        "operationId": "findPetById",
        "produces": [
          "application/json",
          "application/xml",
          "text/xml",
          "text/html"
        ],
        "parameters": [
          {
            "name": "id",
            "in": "path",
            "description": "ID of pet to fetch",
            "required": true,
            "type": "integer",
            "format": "int64"
          }
        ],
        "responses": {
          "200": {
            "description": "pet response",
            "schema": {
              "$ref": "#/definitions/Pet"
            }
          },
          "default": {
            "description": "unexpected error",
            "schema": {
              "$ref": "#/definitions/ErrorModel"
            }
          }
        }
      },
      "delete": {
        "description": "deletes a single pet based on the ID supplied",
        "operationId": "deletePet",
        "parameters": [
          {
            "name": "id",
            "in": "path",
            "description": "ID of pet to delete",
            "required": true,
            "type": "integer",
            "format": "int64"
          }
        ],
        "responses": {
          "204": {
            "description": "pet deleted"
          },
          "default": {
            "description": "unexpected error",
            "schema": {
              "$ref": "#/definitions/ErrorModel"
            }
          }
        }
      }
    }
  },
  "definitions": {
    "Pet": {
      "type": "object",
      "allOf": [
        {
          "$ref": "#/definitions/NewPet"
        },
        {
          "required": [
            "id"
          ],
          "properties": {
            "id": {
              "type": "integer",
              "format": "int64"
            }
          }
        }
      ]
    },
    "NewPet": {
      "type": "object",
      "required": [
        "name"
      ],
      "properties": {
        "name": {
          "type": "string"
        },
        "tag": {
          "type": "string"
        }
      }
    },
    "ErrorModel": {
      "type": "object",
      "required": [
        "code",
        "message"
      ],
      "properties": {
        "code": {
          "type": "integer",
          "format": "int32"
        },
        "message": {
          "type": "string"
        }
      }
    }
  }
}
'''

plantuml_sequence_diagrams = '''@startuml

title ペット
header %page% of %lastpage%
footer Copyright(c) All rights reserved.

autoactivate on
autonumber "<b>[00]"

actor ユーザー as user
entity ペットストアUI as ui
entity ペットストアAPI as api
database データベース as db

== ペットのリストを取得 ==

user -> ui : ペットリストボタンをクリック
ui -> api : findPets
api -> db : SELECT * FROM pets
return
return 200, pet response
return ペットのリストを表示

== ペットの詳細を取得 ==

user -> ui : ペットの詳細ボタンをクリック
ui -> api : findPetById
api -> db : SELECT * FROM pets WHERE id = ${pet_id}
return
return 200, pet response
return ペットの詳細を表示

== ペットを購入 ==

user -> ui : ペットの購入ボタンをクリック
ui -> api : addPet
group transaction
    api -> db : transaction
    api -> db : INSERT INTO orders VALUES (${user_id}, ${pet_id}, ${datetime}, ${created_at}, ${updated_at})
    api -> db : DELETE FROM pets WHERE pet_id = ${pet_id}
    api -> db : commit
    return
    return
    return
    return
end
return 200, pet response
return 購入完了画面を表示

== ペットの登録を削除 ==

user -> ui : ペットの登録削除ボタンをクリック
ui -> api : deletePet
group transaction
    api -> db : DELETE FROM pets WHERE pet_id = (SELECT id FROM orders WHERE pet_id = ${pet_id} AND user_id = ${user_id})
    api -> db : DELETE FROM orders WHERE pet_id = ${pet_id} AND user_id = ${user_id}
    return
    return
end
return 204, pet deleted
return 削除完了画面を表示

== ペットの購入に失敗 ==

user -> ui : ペットの購入ボタンをクリック
ui -> api : addPet
group transaction
    api -> db : transaction
    api -> db : INSERT INTO orders VALUES (${user_id}, ${pet_id}, ${datetime}, ${created_at}, ${updated_at})
    api -> db !! : DELETE FROM pets WHERE pet_id = ${pet_id}
    api -> db : rollback
    return
    return
    return
end
return 500, unexpected error
return 購入失敗画面を表示

@enduml
'''

plantuml_class_diagrams = '''@startuml

class User {
    - id
    - name
}

class Pet {
    - id
    - name
    - type
    - tag
}

class Order {
    - id
    - user_id
    - pet_id
}

Order "1" -- "*" User
Order "1" -- "*" Pet

@enduml
'''

revised_prompt = '''Generate unit test code in Python using the unittest framework and unittest.mock's patch module for given text materials such as OpenAI Specifications (JSON or YAML), PlantUML sequence diagrams (PlantUML format), and PlantUML class diagrams (PlantUML format), following the PEP 8 coding standards.

Key aspects to consider:

* Aim for more than 90% code coverage, with 100% being preferred.
* Consider edge cases and error handling from the given materials step-by-step.
* Follow the AAA (Arrange-Act-Assert) style without including comments in the test code.
* Generate test cases and prepare test data step-by-step, allowing the user to specify test targets/objectives based on the actors of the sequence diagrams.
* Base test cases on class diagrams first, then sequence diagrams, and finally OpenAPI specifications for API servers.
* Define immutable test data in the setUp method of the test class that inherits unittest.TestCase, and specify additional test data in test cases if needed.
* Use the @patch('...') decorator for mocking instead of with patch('...') as mock_foo: statement.
* Apply the spy test double pattern using the assert_called_once_with method of patch.object when applicable, and use the with patch.object('...') as mock_foo: statement in spy test double pattern cases.
* Be mindful of token limitations and avoid exceeding them when generating test code.
* Never apply test doubles like stub, mock, and the rest for the test targets/objectives.
* Ensure that the test code is MECE (mutually exclusive and collectively exhaustive), structured, organized, and clean.
'''

test_target = "ペットストアAPI"

prompt = f'''/*
{revised_prompt}

* Test target: {test_target}

* OpenAPI Specifications:
"""
{openapi_spec}
"""

* PlantUML sequence diagrams:
"""
{plantuml_sequence_diagrams}
"""

* PlantUML class diagrams:
"""
{plantuml_class_diagrams}
"""
*/
'''

messages = [
    {"role":"system","content":"You are an AI assistant that helps people find information."},
    {"role":"user","content":prompt}
]

response = openai.ChatCompletion.create(
  engine="gpt-4-32k-0314",
  messages = messages,
  temperature=0.7,
  max_tokens=1600,
  top_p=0.95,
  frequency_penalty=0,
  presence_penalty=0,
  stop=None)

print(response)

code = response.choices[0].message.content.split('```')[1].lstrip('python\n')
code_string = f'''{code}'''
print(code)

import unittest
from unittest.mock import patch
from petstore_api import PetStoreAPI, Pet, Order


class TestPetStoreAPI(unittest.TestCase):
    def setUp(self):
        self.api = PetStoreAPI()
        self.pets = [
            Pet(1, "Dog", "Bulldog", "friendly"),
            Pet(2, "Cat", "Persian", "lazy"),
            Pet(3, "Bird", "Parrot", "talkative")
        ]
        self.orders = [
            Order(1, 1, 1),
            Order(2, 2, 2)
        ]

    @patch('petstore_api.PetStoreAPI._get_pets_from_db')
    def test_find_pets(self, mock_get_pets_from_db):
        mock_get_pets_from_db.return_value = self.pets
        result = self.api.find_pets()
        self.assertEqual(result, self.pets)
        mock_get_pets_from_db.assert_called_once_with()

    @patch('petstore_api.PetStoreAPI._get_pet_from_db_by_id')
    def test_find_pet_by_id(self, mock_get_pet_from_db_by_id):
        pet_id = 1
        mock_get_pet_from_db_by_id.return_value = self.pets[0]
        result = self.api.find_pet_by_id(pet_id)
        self.assertEqual(result, self.pets[0])
        mock_get_pet_from_db_by_id.assert_called_once_with(pet_id)

    @patch('petstore_api.PetStoreAPI._delete_pet_from_db')
    @patch('petstore_api.PetStoreAPI._delete_order_from_db')
    def test_delete_pet(self, mock_delete_order_from_db, mock_delete_pet_from_db):
        pet_id = 1
        user_id = 1
        self.api.delete_pet(pet_id, user_id)
        mock_delete_order_from_db.assert_called_once_with(pet_id, user_id)
        mock_delete_pet_from_db.assert_called_once_with(pet_id)

    @patch('petstore_api.PetStoreAPI._add_pet_to_db')
    @patch('petstore_api.PetStoreAPI._create_order_in_db')
    def test_add_pet(self, mock_create_order_in_db, mock_add_pet_to_db):
        user_id = 3
        pet_id = 3
        mock_create_order_in_db.return_value = Order(3, user_id, pet_id)
        mock_add_pet_to_db.return_value = self.pets[2]
        result = self.api.add_pet(user_id, pet_id)
        self.assertEqual(result, self.pets[2])
        mock_create_order_in_db.assert_called_once_with(user_id, pet_id)
        mock_add_pet_to_db.assert_called_once_with(pet_id)


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

messages.append(response.choices[0].message)

prompt = "I would like you to add test cases of transaction failure on delete_pet. You should change the test class name to be different from existing test class names."

messages.append({"role":"user","content":prompt})

response = openai.ChatCompletion.create(
  engine="gpt-4-32k-0314",
  messages = messages,
  temperature=0.7,
  max_tokens=1600,
  top_p=0.95,
  frequency_penalty=0,
  presence_penalty=0,
  stop=None)

print(response)

code = response.choices[0].message.content.split('```')[1].lstrip('python\n')
code_string = f'''{code}'''
print(code_string)

import unittest
from unittest.mock import patch
from petstore_api import PetStoreAPI, Pet, Order


class TestPetStoreAPI(unittest.TestCase):
    # ... (previous test cases)

    # Add any new test cases for the existing class here


class TestPetStoreAPITransactionFailures(unittest.TestCase):
    def setUp(self):
        self.api = PetStoreAPI()
        self.pets = [
            Pet(1, "Dog", "Bulldog", "friendly"),
            Pet(2, "Cat", "Persian", "lazy"),
            Pet(3, "Bird", "Parrot", "talkative")
        ]
        self.orders = [
            Order(1, 1, 1),
            Order(2, 2, 2)
        ]

    @patch('petstore_api.PetStoreAPI._delete_pet_from_db')
    @patch('petstore_api.PetStoreAPI._delete_order_from_db')
    def test_delete_pet_transaction_failure(self, mock_delete_order_from_db, mock_delete_pet_from_db):
        pet_id = 1
        user_id = 1
        mock_delete_order_from_db.side_effect = Exception("Transaction failure")

        with self.assertRaises(Exception) as context:
            self.api.delete_pet(pet_id, user_id)

        self.assertEqual(str(context.exception), "Transaction failure")
        mock_delete_order_from_db.assert_called_once_with(pet_id, user_id)
        mock_delete_pet_from_db.assert_not_called()


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