import asyncio
import os
from dataclasses import dataclass, field

from dotenv import load_dotenv
from typing_extensions import Never

from agent_framework import Agent, Executor, WorkflowBuilder, WorkflowContext, handler
from agent_framework.openai import OpenAIChatClient, OpenAIChatOptions

load_dotenv()

_VERIFY_SYSTEM_PROMPT = """\
You are an assistant tasked with determining whether a conversation between a human and a bot \
will continue or not. Your outputs are limited to "[STOP]" or "[CONTINUE]". When you predict \
that the conversation will go on, you should respond with "[CONTINUE]". If you believe the \
conversation has come to an end, respond with "[STOP]".

Examples:

Example 1:
Conversation:
Human: Hey Bot, what's your favorite movie?
Bot: I don't watch movies, but I can help you find information about any movie you like!
Human: Can you tell me about the latest Marvel movie?
Bot: The latest Marvel movie is "Spider-Man: No Way Home". It features Peter Parker dealing \
with the fallout after his identity is revealed. Want to know more about it?
output: [CONTINUE]

Example 2:
Conversation:
Human: Hey Bot, do you know any good Italian restaurants nearby?
Bot: I can't access current location data, but I can suggest looking up Italian restaurants \
on a local review site like Yelp or Google Reviews.
Human: Thanks for the tip. I'll check it out.
Bot: You're welcome! Enjoy your meal. If you need more help, just ask.
output: [STOP]

Instruction:
A conversation is considered to have ended if:
1. The Bot's final response only contains polite expressions without substantive content \
for human to inquire about.
2. In the last round of the conversation, the Human did not ask the Bot any questions."""

_HUMAN_SYSTEM_PROMPT = """\
You are an assistant playing as a random human engaging in a conversation with a digital \
companion, Bot. Your task is to follow the instruction below to role-play as a random human \
in a conversation with Bot, responding to Bot in a manner that a human would say.

Example:
This example illustrates how to generate a conversational response to Bot as a human would:
Conversation:
Human: Bot, what's your favorite movie?
Bot: I don't watch movies, but I can help you find information about any movie you like!
Human: Can you tell me about the latest Marvel movie?
Bot: The latest Marvel movie is "Spider-Man: No Way Home". It features Peter Parker dealing \
with the fallout after his identity is revealed. Want to know more about it?
Human: Yes, can you suggest where I can watch it?

Instruction:
1. Your reply to the Bot should mimic how a human would typically engage in conversation, \
asking questions or making statements that a person would naturally say in response.
2. Do not use interjections.
3. Provide a straightforward, factual response without expressions of surprise, admiration, \
or evaluative comments for Bot's response.
4. Focus on directly asking a question about Bot's response in the last exchange. \
The question should be concise, and without punctuation marks in the middle.
5. Avoid creating any messages that appear to come from the Bot. Your response should not \
contain content that could be mistaken as generated by the Bot, maintaining a clear \
distinction between your input as the Human and the Bot's contributions to the conversation.
6. Your reply should not contain "\\n", this is a reserved character."""


@dataclass
class QuestionSimInput:
    chat_history: list = field(default_factory=list)
    question_count: int = 3


def _format_chat_history(chat_history: list) -> str:
    parts = []
    for item in chat_history:
        parts.append(f"Human: {item['inputs']['question']}")
        parts.append(f"Bot: {item['outputs']['answer']}")
    return "\n".join(parts)


class VerifyAndSimulateExecutor(Executor):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        client = OpenAIChatClient(
            azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
            model=os.environ.get("AZURE_OPENAI_DEPLOYMENT", "gpt-4"),
            api_key=os.environ["AZURE_OPENAI_API_KEY"],
        )
        self._verify_agent = Agent(
            client=client,
            name="VerifyAgent",
            instructions=_VERIFY_SYSTEM_PROMPT,
        )
        self._human_agent = Agent(
            client=client,
            name="HumanSimAgent",
            instructions=_HUMAN_SYSTEM_PROMPT,
        )

    @handler
    async def process(self, sim_input: QuestionSimInput, ctx: WorkflowContext[Never, str]) -> None:
        history_text = _format_chat_history(sim_input.chat_history)

        # Step 1: Verify if conversation should continue
        verify_prompt = (
            f"Read the following conversation and respond:\n"
            f"Conversation:\n{history_text}\noutput:"
        )
        verify_response = await self._verify_agent.run(
            verify_prompt,
            options=OpenAIChatOptions(temperature=0, top_p=1.0),
        )
        stop_or_continue = verify_response.text.strip()

        # Step 2: Check if we should stop
        if "stop" in stop_or_continue.lower():
            await ctx.yield_output("[STOP]")
            return

        # Step 3: Generate human-like questions
        human_prompt = (
            f"Read the following conversation and respond:\n"
            f"Conversation:\n{history_text}\nHuman:"
        )

        # Use the OpenAI client directly for n>1 completions
        from openai import AzureOpenAI

        client = AzureOpenAI(
            api_key=os.environ["AZURE_OPENAI_API_KEY"],
            api_version=os.environ.get("AZURE_OPENAI_API_VERSION", "2024-02-15-preview"),
            azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
        )

        messages = [
            {"role": "system", "content": _HUMAN_SYSTEM_PROMPT},
            {"role": "user", "content": human_prompt},
        ]
        completion = client.chat.completions.create(
            model=os.environ.get("AZURE_OPENAI_DEPLOYMENT", "gpt-4"),
            messages=messages,
            temperature=1.0,
            top_p=1.0,
            presence_penalty=0.8,
            frequency_penalty=0.8,
            n=sim_input.question_count,
            stop=["Human:", "Bot:"],
        )

        questions = []
        for choice in completion.choices:
            response = getattr(choice.message, "content", "")
            if response:
                questions.append(response)

        await ctx.yield_output("\n".join(questions))


def create_workflow():
    """Create a fresh workflow instance.

    MAF workflows do not support concurrent execution, so each
    concurrent caller needs its own workflow instance.
    """
    _executor = VerifyAndSimulateExecutor(id="verify_and_simulate")
    return (
        WorkflowBuilder(name="QuestionSimulationWorkflow", start_executor=_executor)
        .build()
    )


async def main():
    workflow = create_workflow()
    chat_history = [
        {
            "inputs": {"question": "Can you introduce something about large language model?"},
            "outputs": {
                "answer": (
                    "A large language model (LLM) is a type of language model that is distinguished "
                    "by its ability to perform general-purpose language generation and understanding. "
                    "These models learn statistical relationships from text documents through a "
                    "self-supervised and semi-supervised training process that is computationally intensive."
                ),
            },
        }
    ]
    result = await workflow.run(
        QuestionSimInput(chat_history=chat_history, question_count=3)
    )
    print(f"Output:\n{result.get_outputs()[0]}")


if __name__ == "__main__":
    asyncio.run(main())