import semantic_kernel as sk
from semantic_kernel.connectors.ai.open_ai import OpenAIChatCompletion
from dotenv import load_dotenv, find_dotenv
import os

import random,asyncio
from semantic_kernel import KernelContext
from semantic_kernel.plugin_definition import kernel_function, kernel_function_context_parameter




load_dotenv(find_dotenv())




api_key = os.getenv('OPENAI_API_KEY')
model = OpenAIChatCompletion(
    "gpt-3.5-turbo-1106",
    api_key
)


kernel = sk.Kernel()



kernel.add_text_completion_service("gpt-3.5", model)



class GenerateNumberPlugin:
    @kernel_function(
        description="Generate a random number between min and max",
        name="GenerateNumber",
    )
    @kernel_function_context_parameter(name="min", description="Minimum number of paragraphs.")
    @kernel_function_context_parameter(name="max", description="Maximum number of paragraphs.", default_value="10")
    def generate_number(self, context: KernelContext) -> str:
        try:
            return str(random.randint(int(context["min"]), int(context["max"])))
        except ValueError as e:
            print(f"Invalid input {context['min']} {context['max']}")
            raise e

class GenerateNamesPlugin:
    @kernel_function(description="Generate character names", name="generate_names")
    def generate_names(self) -> str:
        """
        Generate two names.
        Returns:
            str
        """
        names = {"Hoagie", "Hamilton", "Bacon", "Pizza", "Boots", "Shorts", "Tuna"}
        first_name = random.choice(list(names))
        names.remove(first_name)
        second_name = random.choice(list(names))
        return f"{first_name}, {second_name}"
    

generate_number_plugin = kernel.import_plugin(GenerateNumberPlugin(), "GenerateNumberPlugin")
generate_number = generate_number_plugin["GenerateNumber"]


generate_names_plugin = kernel.import_plugin(GenerateNamesPlugin(), plugin_name="GenerateNames")
generate_names = generate_names_plugin["generate_names"]


sk_prompt = """
Write a short story about two Corgis on an adventure.
The story must be:
- G rated
- Have a positive message
- No sexism, racism or other bias/bigotry
- Be exactly {{$paragraph_count}} paragraphs long
- Be written in this language: {{$language}}
- The two names of the corgis are {{GenerateNames.generate_names}}
"""

print(sk_prompt)

corgi_story = kernel.create_semantic_function(
    prompt_template=sk_prompt,
    function_name="CorgiStory",
    plugin_name="CorgiPlugin",
    description="Write a short story about two Corgis on an adventure.",
    max_tokens=500,
    temperature=0.5,
    top_p=0.5,
)


context_variables = sk.ContextVariables(variables={"min": "2", "max": "5", "language": "中文"})

num=asyncio.run(generate_number.invoke(variables=context_variables))
print(num.result)
context_variables["paragraph_count"] = str(num.result)

names =  asyncio.run(generate_names.invoke())
print(names)

story  =  asyncio.run(corgi_story.invoke(variables=context_variables))
print(story)