from langchain_openai import ChatOpenAI
from openai import OpenAI
import logging
from asyncio import gather
import os
import sys

from tools import openai_url_invoke
from agent_protocol import MultiAgentState

module_paths = ["../"]
file_path = os.path.dirname(__file__)
os.chdir(file_path)

for module_path in module_paths:
    full_path = os.path.normpath(os.path.join(file_path, module_path))
    sys.path.append(full_path)

from config import gconfig


# --------------------------------------------------------------------------------------------
# Config -------------------------------------------------------------------------------------
# --------------------------------------------------------------------------------------------


# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Model IDs
model_id_c37 = "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
model_id_c35 = "us.anthropic.claude-3-5-haiku-20241022-v1:0"
model_id_nova = "us.amazon.nova-lite-v1:0"
model_id_llama = "meta.llama3-3-70b-instruct-v1:0"

config = {
    'model_id': model_id_nova,
    'internal_file_path': '',
    'tools': [],
    'master_server_client': None,
    'dispatcher_system_message': '',
}


def set_agent_config(options: dict):
    for key in options:
        config[key] = options[key]
    logging.info(f'Agent config set to {config}')


# OpenAI chat client
openai_client = OpenAI()


# --------------------------------------------------------------------------------------------
# Deployment Nodes ---------------------------------------------------------------------------
# --------------------------------------------------------------------------------------------


async def tools_selector_node(state: MultiAgentState):
    logging.info(f'Selecting tools from {[tool['function']['name'] for tool in config.get('tools')]}...')

    # Select tools
    messages = [
        {"role": "system", "content": config.get('dispatcher_system_message')},
        {"role": "user", "content": state.get('question')},
    ]

    response = openai_client.chat.completions.create(
        model=config.get('model_id'),
        max_tokens=3000,
        messages=messages,
        tools=config.get('tools')
    )

    tool_calls = response.choices[0].message.tool_calls
    logging.info(f'Selected tools: {tool_calls}')

    # Call selected tools
    if tool_calls:
        tool_callables = [
            config.get('master_server_client').call_tool(
                tool.function.name, eval(tool.function.arguments)
            )
            for tool in tool_calls
        ]

        results = await gather(*tool_callables)
        external_data = [result.content[0].text for result in results]

        logging.info(f'Tool results: {external_data}')
        return {'tools_requested': tool_calls, 'external_data': str(external_data)}

    # Safety in case the model chooses to generate its own response
    response = response.choices[0].message.content
    logging.info(f'Model-generated response: {response}')
    return {'tools_requested': None, 'external_data': response}


def judge_node(state: MultiAgentState):
    logging.info(f'Judging response...')
    # Answers:
    # GOOD - send to client
    # BAD - try again

    decision, _ = openai_url_invoke(
        gconfig.get('judge_model_id'),
        state.get('question'),
        str(state.get('external_data')),
        f'You are a seasoned expert. Your role is to determine the quality of the answer generated by '
        f'a team of other AI models. Output "GOOD" if the response answers the question and matches '
        f'real world consensus, otherwise output "BAD". Do not output anything else besides those two '
        f'options. Output in all capital letters.',
        gconfig.get('judge_model_service_url'),
    )
    decision = decision.strip().upper()

    logging.info(f'Judge decision: {decision}')
    return {'qa_assessment': decision}


def judge_decision(state: MultiAgentState):
    logging.info(f'Routing judge decision...')

    return state.get('qa_assessment')
