from modelscope_studio.components.pro.chatbot import (
    ChatbotActionConfig,
    ChatbotBotConfig,
    ChatbotUserConfig,
    ChatbotWelcomeConfig,
)

MIN_TEMPERATURE = 0
MAX_TEMPERATURE = 1
DEFAULT_TEMPERATURE = 0.6

MIN_MAX_NEW_TOKENS = 4096
MAX_MAX_NEW_TOKENS = 65536
DEFAULT_MAX_NEW_TOKENS = 65536

MIN_TOP_P = 0
MAX_TOP_P = 1
DEFAULT_TOP_P = 0.95


DEFAULT_MODEL = "tiiuae/Falcon-H1R-7B"
MODEL_OPTIONS = [
    {
        "label": "Falcon-H1R-7B",
        "modelId": "tiiuae/Falcon-H1R-7B",
        "value": "tiiuae/Falcon-H1R-7B",
        "link": "https://huggingface.co/tiiuae/Falcon-H1R-7B",
        "api_key_env": "X_API_KEY",
        "base_url_env": "X_URL",
    }
]
MODEL_OPTIONS_MAP = {model["value"]: model for model in MODEL_OPTIONS}

DEFAULT_LOCALE = "en_US"

DEFAULT_THEME = {
    "token": {
        "colorPrimary": "#7422f2",
    }
}

DEFAULT_SETTINGS = {
    "model": DEFAULT_MODEL,
    "temperature": DEFAULT_TEMPERATURE,
    "max_new_tokens": DEFAULT_MAX_NEW_TOKENS,
    "top_p": DEFAULT_TOP_P,
}


# Save history in browser
save_history = True

CHAR_YIELD_SIZE = 100  # Update UI every 100 characters
KEPT_TURNS = 10


# Chatbot Config
def user_config(disabled_actions=None):
    return ChatbotUserConfig(
        class_names=dict(content="user-message-content"),
        actions=[
            "copy",
            "edit",
            ChatbotActionConfig(
                action="delete",
                popconfirm=dict(
                    title="Delete the message",
                    description="Are you sure to delete this message?",
                    okButtonProps=dict(danger=True),
                ),
            ),
        ],
        disabled_actions=disabled_actions,
    )


def bot_config(disabled_actions=None):
    return ChatbotBotConfig(
        actions=[
            "copy",
            "edit",
            ChatbotActionConfig(
                action="retry",
                popconfirm=dict(
                    title="Re-generate the response",
                    description="Re-generate the response will also delete all subsequent messages.",
                    okButtonProps=dict(danger=False),
                ),
            ),
            ChatbotActionConfig(
                action="delete",
                popconfirm=dict(
                    title="Delete the response",
                    description="Are you sure to delete this response?",
                    okButtonProps=dict(danger=True),
                ),
            ),
        ],
        avatar="./photos/falcon_logo.png",
        disabled_actions=disabled_actions,
    )


def welcome_config():
    return ChatbotWelcomeConfig(
        # variant="filled",
        variant="borderless",
        icon="./photos/falcon_logo.png",
        title="Falcon-H1R",
        description="Select your preferred reasoning model size and set your configs on the 'Settings' menu and enter your message to start chatting.",
        prompts=dict(
            styles={
                "list": {
                    "width": "100%",
                },
                "item": {
                    "flex": 1,
                },
            },
            items=[
                {
                    "label": "📟 Write Python code",
                    "children": [
                        {
                            "description": "Write a python function that returns the factorial of a number using recursion.",
                        },
                    ],
                },
                {
                    "label": "🔢 Solve math questions",
                    "children": [
                        {
                            "description": "Solve 3x^2 + 2x + 1 = 0",
                        },
                    ],
                },
                {
                    "label": "🧠 Logical reasoning",
                    "children": [
                        {
                            "description": "A train leaves Station A at 9:00 AM traveling at 60 mph toward Station B. Another train leaves Station B at 10:00 AM traveling at 80 mph toward Station A. If the stations are 280 miles apart, at what time do the trains meet?",
                        },
                    ],
                },
            ],
        ),
    )