# Import FastAPI Libs
from pydantic import BaseModel

# Import Dashscope Service Libs
from http import HTTPStatus
from dashscope import Generation
from dashscope.api_entities.dashscope_response import Role
import dashscope
import json
import os

# Import your custom logging manager
import logging
from common.script.logging_manager import LoggingManager

# Initialize logging manager
curr_module = os.path.basename(__file__)
lm = LoggingManager.get_instance()

# ------------------------------------------------
# Define constants
# ------------------------------------------------
c_ret_code_success = 0
c_ret_code_error = 1

gen_ai_id = 'aliyun'
credential_file = '/home/oracle/python/apex/credential/credential.json'

# ------------------------------------------------
# Get credentials
# ------------------------------------------------
def get_credentials(file: str, gen_ai_id: str) -> str:

    try:
        # Open and read the JSON file
        with open(file, 'r') as credential_file:
            data = json.load(credential_file)

        # Access data from the JSON
        api_key = data['gen-ai'][gen_ai_id]['api_key']

        return api_key

    except Exception as e:

        lm.log(logging.ERROR, "Exception occurred in get_credentials: %s", str(e), exc_info=True)
        return None


# ------------------------------------------------
# Define GenAI Message Service Request and Response
# ------------------------------------------------
class Massage(BaseModel):
    role: str
    text: str

class send_message_req_basemodel(BaseModel):
    model: str
    messages: list[Massage]

class send_message_res_basemodel(BaseModel):
    result: int
    result_message: str
    message_res: str


# ------------------------------------------------
# Function to call dashscope SDK
# ------------------------------------------------
def send_message(messages_req: send_message_req_basemodel) -> send_message_res_basemodel:

    result = c_ret_code_success
    result_message = ""

    try:

        # Get API Key
        api_key = get_credentials(credential_file, gen_ai_id)
        if (api_key == None):
            result = c_ret_code_error
            # Set return tuple for normal result
            send_message_res = send_message_res_basemodel(
                result          = result,
                result_message  = 'Failed to get API key',
                message_res     = ''
            )
            return send_message_res
        else:
            dashscope.api_key = api_key

        # Set system prompt
        messages = [{'role': Role.SYSTEM, 'content': 'You are a helpful assistant.'}]

        # Loop input message list
        for message in messages_req.messages:
            # Set user and assistant prompt
            messages.append({'role': message.role, 'content': message.text})

        lm.log(logging.INFO, curr_module, 'message: ', str(messages))

        # Get Model
        # model = Models()
        match messages_req.model:
            case 'qwen-max':
                model = Generation.Models.qwen_max
            case 'qwen-turbo':
                model = Generation.Models.qwen_turbo
            case 'qwen-plus':
                model = Generation.Models.qwen_plus

        # Use SDK to get answer
        response = Generation.call(
            model,                  # Generation.Models.qwen_turbo,
            messages=messages,
            result_format='message',  # set the result to be 'message' format.
        )
        if response.status_code == HTTPStatus.OK:

            lm.log(logging.INFO, curr_module, 'Normal response', str(response))

            send_message_res = send_message_res_basemodel(
                result = result,
                result_message = '',
                message_res = response['output']['choices'][0]['message']['content']
            )
            return send_message_res

        else:

            lm.log(logging.INFO, curr_module, 'Error response', str(response))

            send_message_res = send_message_res_basemodel(
                result = result,
                result_message = '',
                message_res = response.message
            )
            return send_message_res

    except Exception as e:

        lm.log(logging.ERROR, "Exception occurred: %s", str(e), exc_info=True)

        result = c_ret_code_error

        # Set return tuple for abnormal result
        send_message_res = send_message_res_basemodel(
            result          = result,
            result_message  = str(e),
            message_res     = ''
        )

        return send_message_res