import json
import os
from openai import OpenAI

class Orchestrator:
    def __init__(self, kb_path="hf_mcp_agent/knowledge_base.json"):
        with open(kb_path, "r") as f:
            self.kb = json.load(f)
        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    def get_agent_for_intent(self, user_input):
        # Use a fast call to classify intent based on mapping
        prompt = f"Classify the following user input into one of these categories: {list(self.kb['routing_logic']['intent_mapping'].keys())}. 
Input: {user_input}
Output just the category name."
        response = self.client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}]
        )
        intent = response.choices[0].message.content.strip().lower()
        agent_key = self.kb['routing_logic']['intent_mapping'].get(intent, self.kb['routing_logic']['default_agent'])
        return agent_key

    def get_agent_config(self, agent_key):
        agent = self.kb['agents'][agent_key]
        # Map string tool names to actual function definitions from KB
        tools_definition = []
        for tool_name in agent['tools']:
            if tool_name in self.kb['tool_definitions']:
                tools_definition.append({
                    "type": "function",
                    "function": {
                        "name": tool_name,
                        "description": self.kb['tool_definitions'][tool_name]['description'],
                        "parameters": {"type": "object", "properties": self.kb['tool_definitions'][tool_name]['parameters']}
                    }
                })
        return agent['system_prompt'], tools_definition