import asyncio

import writer as wf
from dotenv import load_dotenv
from writer import WriterState
from writerai import AsyncWriter
from writerai.types.chat_chat_params import Message

load_dotenv()

def handle_user_input(payload: str,  state: WriterState) -> None:
    asyncio.run(_ask_models(payload, state))


def handle_prompt_button(context: dict, state: WriterState) -> None:
    state["user-prompt"] = state[context["target"]+'-full']
    asyncio.run(_ask_models(state[context["target"]+'-full'], state))


async def _ask_models(prompt: str, state: WriterState) -> None:
    state["palmyra-x-004-conversation"].append(Message(role="user", content=prompt))
    state["palmyra-creative-conversation"].append(Message(role="user", content=prompt))

    async_writer_client = AsyncWriter()

    await asyncio.gather(
        _perform_async_streaming(async_writer_client, "palmyra-x-004", state),
        _perform_async_streaming(async_writer_client, "palmyra-creative", state),
    )


async def _perform_async_streaming(client: AsyncWriter, model: str, state: WriterState) -> None:
    try:
        response = await client.chat.chat(
            model=model,
            messages=state[f"{model}" + "-conversation"],
            stream=True,
            max_tokens=16384
        )

        response_message = ""
        state[f"{model}" + "-response"] = ""

        async for message in response:
            content = message.choices[0].delta.content
            content = content if content is not None else ""
            response_message += content
            state[f"{model}" + "-response"] += content

        state[f"{model}" + "-conversation"].append(Message(role="assistant", content=response_message))

    except Exception as e:
        response_message = "Something went wrong. Please, try again."
        state[f"{model}" + "-conversation"].append(Message(role="assistant", content=response_message))
        print(e)

initial_state = wf.init_state({
    "palmyra-x-004-conversation": [],
    "palmyra-x-004-response": "Model response will appear here...",
    "palmyra-creative-conversation": [],
    "palmyra-creative-response": "Model response will appear here...",
    "user-prompt": "",
    "prompt-left-button": "Brainstorm bakery strategies",
    "prompt-left-button-full": "Imagine you're a struggling small-town bakery competing with a chain that opened across the street. Brainstorm unconventional strategies to win over customers without lowering prices.",
    "prompt-center-button": "Explain AI to a high schooler",
    "prompt-center-button-full": "Write a guide for a programmer who wants to explain their AI side project to a high schooler. The explanation must be engaging, simple, and use humorous analogies, while avoiding technical jargon.",
    "prompt-right-button": "Zero gravity game",
    "prompt-right-button-full": "Design a game that could only exist in zero gravity."
})

initial_state.import_stylesheet("style", "/static/custom.css")