import os
from openai import OpenAI
import logging
from asyncio import gather
from pydantic import BaseModel
from typing import Any

from .tools import openai_url_invoke
from .agent_protocol import MultiAgentState
from mcp_master.config import global_config as gconfig
from mcp_master.config import ConfigError

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


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


class Config(BaseModel):
    model_id: str = '',
    tools: list = [],
    master_server_client: Any = None,
    dispatcher_system_message: str = '',


config = Config()

# OpenAI chat client
openai_client = None


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


async def tools_selector_node(state: MultiAgentState):
    global openai_client

    logging.info(f'Selecting tools from {[tool['function']['name'] for tool in config.tools]}...')

    # Exit for invalid selector_model_id values
    config.model_id = gconfig.selector_model_id
    if config.model_id is None or len(config.model_id) == 0:
        raise ConfigError(f'Ensure your selector_model_id is properly configured via set_config(). It is currently {config.model_id}.')

    # Exit for invalid API key
    if gconfig.OPENAI_API_KEY is None or len(gconfig.OPENAI_API_KEY) == 0:
        raise ConfigError(f'Ensure your OPENAI_API_KEY is properly configured via set_config(). It is currently {gconfig.OPENAI_API_KEY}.')
    os.environ["OPENAI_API_KEY"] = gconfig.OPENAI_API_KEY

    # Exit for empty API key
    if gconfig.OPENAI_BASE_URL is not None:
        if len(gconfig.OPENAI_BASE_URL) == 0:
            raise ConfigError(
                f'Ensure your OPENAI_BASE_URL is properly configured via set_config(). It is currently empty.')

        # Save API key to environment variables
        os.environ["OPENAI_BASE_URL"] = gconfig.OPENAI_BASE_URL

    # Initialize openAI chat client
    if openai_client is None:
        openai_client = OpenAI()

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

    response = openai_client.chat.completions.create(
        model=config.model_id,
        max_tokens=3000,
        messages=messages,
        tools=config.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.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.judge_model_id,
        state.question,
        str(state.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.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.qa_assessment
