import json
import sys

from pygeai import logger
from pygeai.chat.clients import ChatClient
from pygeai.chat.iris import Iris
from pygeai.cli.commands import Command, Option, ArgumentsEnum
from pygeai.cli.commands.builders import build_help_text
from pygeai.cli.commands.common import get_messages, get_boolean_value, get_penalty_float_value
from pygeai.cli.texts.help import CHAT_HELP_TEXT
from pygeai.core.common.exceptions import MissingRequirementException, WrongArgumentError
from prompt_toolkit import PromptSession
from prompt_toolkit.history import InMemoryHistory

from pygeai.core.utils.console import Console


def show_help():
    """
    Displays help text in stdout
    """
    help_text = build_help_text(chat_commands, CHAT_HELP_TEXT)
    Console.write_stdout(help_text)


def get_chat_completion(option_list: list):
    model = None
    message_list = []
    stream = False
    temperature = None
    max_tokens = None
    thread_id = None
    frequency_penalty = None
    presence_penalty = None
    top_p = None
    stop = None
    response_format = None
    tools = None
    tool_choice = None
    logprobs = None
    top_logprobs = None
    seed = None
    stream_options = None
    store = None
    metadata = None
    user = None

    for option_flag, option_arg in option_list:
        if option_flag.name == "model":
            model = option_arg
        if option_flag.name == "messages":
            try:
                message_json = json.loads(option_arg)
                if isinstance(message_json, list):
                    message_list = message_json
                elif isinstance(message_json, dict):
                    message_list.append(message_json)
            except Exception as e:
                raise WrongArgumentError(
                    "Each message must be in json format: '{\"role\": \"user\", \"content\": \"message content\"}' "
                    "It can be a dictionary or a list of dictionaries. Each dictionary must contain role and content")
        if option_flag.name == "stream":
            if option_arg:
                stream = get_boolean_value(option_arg)
        if option_flag.name == "temperature":
            temperature = float(option_arg) if option_arg is not None else None
        if option_flag.name == "max_tokens":
            max_tokens = int(option_arg) if option_arg is not None else None
        if option_flag.name == "thread_id":
            thread_id = option_arg
        if option_flag.name == "frequency_penalty":
            if option_arg:
                frequency_penalty = get_penalty_float_value(option_arg)
        if option_flag.name == "presence_penalty":
            if option_arg:
                presence_penalty = get_penalty_float_value(option_arg)
        if option_flag.name == "top_p":
            top_p = float(option_arg) if option_arg else None
        if option_flag.name == "stop":
            try:
                stop = json.loads(option_arg) if option_arg else None
            except json.JSONDecodeError:
                stop = option_arg
        if option_flag.name == "response_format":
            try:
                response_format = json.loads(option_arg) if option_arg else None
            except json.JSONDecodeError:
                raise WrongArgumentError("response_format must be a valid JSON object")
        if option_flag.name == "tools":
            try:
                tools = json.loads(option_arg) if option_arg else None
            except json.JSONDecodeError:
                raise WrongArgumentError("tools must be a valid JSON array")
        if option_flag.name == "tool_choice":
            try:
                tool_choice = json.loads(option_arg) if option_arg else None
            except json.JSONDecodeError:
                tool_choice = option_arg
        if option_flag.name == "logprobs":
            logprobs = get_boolean_value(option_arg) if option_arg else None
        if option_flag.name == "top_logprobs":
            top_logprobs = int(option_arg) if option_arg else None
        if option_flag.name == "seed":
            seed = int(option_arg) if option_arg else None
        if option_flag.name == "stream_options":
            try:
                stream_options = json.loads(option_arg) if option_arg else None
            except json.JSONDecodeError:
                raise WrongArgumentError("stream_options must be a valid JSON object")
        if option_flag.name == "store":
            store = get_boolean_value(option_arg) if option_arg else None
        if option_flag.name == "metadata":
            try:
                metadata = json.loads(option_arg) if option_arg else None
            except json.JSONDecodeError:
                raise WrongArgumentError("metadata must be a valid JSON object")
        if option_flag.name == "user":
            user = option_arg

    messages = get_messages(message_list)

    if not (model and messages):
        raise MissingRequirementException("Cannot perform chat completion without specifying model and messages")

    client = ChatClient()
    result = client.chat_completion(
        model=model,
        messages=messages,
        stream=stream,
        temperature=temperature,
        max_tokens=max_tokens,
        thread_id=thread_id,
        frequency_penalty=frequency_penalty,
        presence_penalty=presence_penalty,
        top_p=top_p,
        stop=stop,
        response_format=response_format,
        tools=tools,
        tool_choice=tool_choice,
        logprobs=logprobs,
        top_logprobs=top_logprobs,
        seed=seed,
        stream_options=stream_options,
        store=store,
        metadata=metadata,
        user=user
    )
    if stream:
        Console.write_stdout("Streaming chat completion:")
        for chunk in result:
            Console.write_stdout(f"{chunk}", end="")
            sys.stdout.flush()
        Console.write_stdout()
    else:
        Console.write_stdout(f"Chat completion detail: \n{result}\n")


chat_completion_options = [
    Option(
        "model",
        ["--model", "-m"],
        "The model needs to address the assistant type and name or bot_id, depending on the Type. Then, the parameters"
        " will vary depending on the type. Its format is as follows: \n"
        "\t\"model\": \"saia:<assistant_type>:<assistant_name>|<bot_id>\"",
        True
    ),
    Option(
        "messages",
        ["--messages", "--msg"],
        "The messages element defines the desired messages to be added. The minimal value needs to be the following, "
        "where the content details the user input.\n"
        "\t{ \n"
        "\t\t\"role\": \"string\", /* user, system and may support others depending on the selected model */ \n"
        "\t\t\"content\": \"string\" \n"
        "\t}\n",
        True
    ),
    Option(
        "stream",
        ["--stream"],
        "If response should be streamed. Possible values: 0: OFF; 1: ON",
        True
    ),
    Option(
        "temperature",
        ["--temperature", "--temp"],
        "Float value to set volatility of the assistant's answers (between 0 and 2)",
        True
    ),
    Option(
        "max_tokens",
        ["--max-tokens"],
        "Integer value to set max tokens to use",
        True
    ),
    Option(
        "thread_id",
        ["--thread-id"],
        "Optional UUID for conversation identifier",
        True
    ),
    Option(
        "frequency_penalty",
        ["--frequency-penalty"],
        "Optional number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency "
        "in the text so far, decreasing the model's likelihood to repeat the same line verbatim.",
        True
    ),
    Option(
        "presence_penalty",
        ["--presence-penalty"],
        "Optional number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in "
        "the text so far, increasing the model's likelihood to talk about new topics.",
        True
    ),
    Option(
        "top_p",
        ["--top-p"],
        "Optional float value for nucleus sampling, where the model considers tokens with top_p probability mass "
        "(between 0 and 1). An alternative to temperature.",
        True
    ),
    Option(
        "stop",
        ["--stop"],
        "Optional string or JSON array of up to 4 sequences where the API will stop generating further tokens.",
        True
    ),
    Option(
        "response_format",
        ["--response-format"],
        "Optional JSON object specifying the output format, e.g., {\"type\": \"json_schema\", \"json_schema\": {...}} "
        "for structured outputs.",
        True
    ),
    Option(
        "tools",
        ["--tools"],
        "Optional JSON array of tools (e.g., functions) the model may call.",
        True
    ),
    Option(
        "tool_choice",
        ["--tool-choice"],
        "Optional string (e.g., \"none\", \"auto\") or JSON object to control which tool is called.",
        True
    ),
    Option(
        "logprobs",
        ["--logprobs"],
        "Optional boolean to return log probabilities of output tokens. Possible values: 0: OFF; 1: ON",
        True
    ),
    Option(
        "top_logprobs",
        ["--top-logprobs"],
        "Optional integer (0-20) specifying the number of most likely tokens to return with log probabilities.",
        True
    ),
    Option(
        "seed",
        ["--seed"],
        "Optional integer for deterministic sampling (in Beta).",
        True
    ),
    Option(
        "stream_options",
        ["--stream-options"],
        "Optional JSON object for streaming options, e.g., {\"include_usage\": true}.",
        True
    ),
    Option(
        "store",
        ["--store"],
        "Optional boolean to store the output for model distillation or evals. Possible values: 0: OFF; 1: ON",
        True
    ),
    Option(
        "metadata",
        ["--metadata"],
        "Optional JSON object with up to 16 key-value pairs to attach to the object.",
        True
    ),
    Option(
        "user",
        ["--user"],
        "Optional string identifier for the end-user to monitor abuse.",
        True
    ),
]


def chat_with_iris():
    iris = Iris()
    messages = list()

    history = InMemoryHistory()
    session = PromptSession(">> Ask Iris: ", history=history)

    Console.write_stdout(f"#=================================================#")
    Console.write_stdout(f"#--------------------- IRIS ----------------------#")
    Console.write_stdout(f"#=================================================#")
    Console.write_stdout(f"# This is the start of your conversation with Iris. Type 'exit' or press Ctrl+C to close the chat.\n")
    Console.write_stdout("""- Iris: Hello! I'm Iris. I'll guide you step by step to create your agent.
First, we need to define some key details for your agent. You can specify its role and purpose or give it a name, and I'll help you set up the rest. Once we have that, we'll refine its knowledge and behavior.""")
    try:
        while (user_input := session.prompt()) != "exit":
            Console.write_stdout(f"- User: {user_input}")
            new_message = {
                "role": "user",
                "content": user_input
            }
            messages.append(new_message)

            result = iris.stream_answer(messages)
            answer = ""
            Console.write_stdout("- Iris: ")
            for chunk in result:
                answer += chunk
                Console.write_stdout(f"{chunk}", end="")
                sys.stdout.flush()
            Console.write_stdout()

            new_answer = {
                "role": "agent",
                "content": answer
            }
            messages.append(new_answer)
    except KeyboardInterrupt:
        print("\nExiting chat...")
        sys.exit(0)

    except Exception as e:
        logger.error(f"Error chatting with Iris: {e}")
        Console.write_stderr("An unexpected error has occurred. Please contact the developers.")


chat_commands = [
    Command(
        "help",
        ["help", "h"],
        "Display help text",
        show_help,
        ArgumentsEnum.NOT_AVAILABLE,
        [],
        []
    ),
    Command(
        "completion",
        ["completion", "comp"],
        "Get chat completion",
        get_chat_completion,
        ArgumentsEnum.REQUIRED,
        [],
        chat_completion_options
    ),
    Command(
        "iris",
        ["iris"],
        "Interactive chat with Iris",
        chat_with_iris,
        ArgumentsEnum.NOT_AVAILABLE,
        [],
        []
    ),

]
