from openai import AsyncOpenAI
import chainlit as cl
from chainlit.input_widget import Select, Switch, Slider, TextInput
import httpx
import os

import utils

BASE_URL=os.getenv("OPENAI_API_BASE")
API_KEY=os.getenv("OPENAI_API_KEY")
PROXY_URL=os.getenv("PROXY_API_BASE") or BASE_URL

client = AsyncOpenAI(
    base_url=PROXY_URL,
    api_key=API_KEY,
)

# Instrument the OpenAI client
cl.instrument_openai()

router={
    "白痴模式": "llm",
    # "疯狗模式": "llm",
    "简单模式": "llm",
}

def get_chat_settings():
    profile = cl.user_session.get("chat_profile")
    print('profile:', profile)

    if profile == "白痴模式":
        return [
            Select(
                id="model",
                label="Model",
                values=utils.Models,
                initial_index=0,
            ),
        ]
    # elif profile == "疯狗模式":
    #     return [
    #         Select(
    #             id="model",
    #             label="Model",
    #             values=utils.Models,
    #             initial_index=0,
    #         ),
    #     ]
    elif profile == "简单模式":
        return [
            Select(
                id="model",
                label="Model",
                values=utils.Models,
                initial_index=0,
            ),
            Slider(
                id="temperature",
                label="Temperature",
                initial=1,
                min=0,
                max=2,
                step=0.1,
            ),
            Slider(
                id="top_p",
                label="Top P",
                initial=1.0,
                min=0.1,
                max=1.0,
                step=0.1,
            ),
            TextInput(
                id="system_prompt",
                label="System Prompt",
                placeholder="You are a helpful assistant.",
            ),
        ]

@cl.password_auth_callback
def auth(username: str, password: str):
    if (username, password) == ("dog", "123"):
        return cl.User(
            identifier="admin",
            display_name="神仙",
            metadata={"role":"master","provider":"credentials","image":utils.Images['admin']},
        )
    elif password == "jj88":
        return cl.User(
            identifier="dog",
            display_name="狗",
            metadata={"role":"dog","provider":"credentials"},
        )
    else:
        return None

@cl.set_chat_profiles
async def chat_profile(user: cl.User):
    print('user:', user)
    return [
        cl.ChatProfile(
            name="白痴模式",
            markdown_description="所有参数默认",
            icon="https://image.acg.lol/file/2025/02/18/image.jpg",
        ),
        cl.ChatProfile(
            name="简单模式",
            markdown_description="可以配置简单的参数",
            icon="https://image.acg.lol/file/2025/02/18/image.jpg",
        ),
    ]

    if user.metadata['role'] == "master":
        return profiles + [
            cl.ChatProfile(
                name="Flux",
                markdown_description="hahah222",
                icon="https://image.acg.lol/file/2025/02/18/image.jpg",
            ),
        ]

    print(profiles)
    return profiles

# @cl.set_starters
# async def starters():
#     return [
#         cl.Starter(
#             label="叫鬼啊",
#             message="hello",
#         ),
#     ]

@cl.on_chat_start
async def start():
    profile = cl.user_session.get("chat_profile")
    print('profile::', profile)
    settings = await cl.ChatSettings(get_chat_settings()).send()


@cl.on_settings_update
async def setup_agent(settings):
    print("on_settings_update", settings)

@cl.on_message
async def on_message(message: cl.Message):
    # params = cl.user_session.get('chat_settings') | {
    #     "extra_headers": {
    #         "x-api-endpoint": BASE_URL,
    #     },
    # }
    # print(params)
    # context = cl.chat_context.to_openai()
    # print(context)
    # response = await client.chat.completions.create(
    #     messages=context,
    #     stream=True,
    #     **params,
    # )
    # msg = cl.Message(content="")
    # async for chunk in response:
    #     await msg.stream_token(chunk.choices[0].delta.content or "")
    #     # await cl.Message(content=response.choices[0].message.content).send()
    # await msg.send()
    profile = cl.user_session.get("chat_profile")
    mode = router[profile]

    if mode == "llm":
        # if profile == "疯狗模式":
        #     return await on_llm_chat(message, "https://api.groq.com/openai/v1")
        return await on_llm_chat(message)
    else:
        return None



async def on_llm_chat(message: cl.Message, base_url=None):
    params = cl.user_session.get('chat_settings') | {
        "extra_headers": {
            "x-api-endpoint": base_url or BASE_URL,
        },
    }
    print(params)
    context = cl.chat_context.to_openai()
    if context[0]['role'] != "system":
        pass
    print(context)
    response = await client.chat.completions.create(
        messages=context,
        stream=True,
        **params,
    )
    msg = cl.Message(content="")
    async for chunk in response:
        await msg.stream_token(chunk.choices[0].delta.content or "")

    await msg.send()