import random

from langchain import PromptTemplate, LLMChain
from langchain.chat_models import ChatOpenAI

from experts.base_expert import BaseExpert


class Conductor(BaseExpert):

    ROLE_DESCRIPTION='''you will take on the role of the conductor for a multi-expert system.'''
    FORWARD_TASK = '''Now, you are presented with an operational optimization-related problem: 
{problem_description}

In this multi-expert system, there are many experts, each of whom is responsible for solving part of the problem.
Your task is to CHOOSE THE NEXT EXPERT TO CONSULT.

The names of the experts and their capabilities are listed below:
{experts_info} 

Experts that have already been commented include: 
{commented_experts}

Please select an expert to consult from the remaining expert names {remaining_experts}.

Please note that the maximum number of asked experts is {max_collaborate_nums}, and you can ask {remaining_collaborate_nums} more times.

You should output the name of expert directly. The next expert is:'''

    def __init__(self, model):
        super().__init__(
            name='Conductor',
            description='An special expert that collaborates all other experts.',
            model=model
        )

        self.llm = ChatOpenAI(
            model_name=model,
            temperature=0,
            max_tokens=50,  # Limit response length for expert selection
            api_key="/",  # Dummy key for RITS
            base_url="https://inference-3scale-apicast-production.apps.rits.fmaas.res.ibm.com/deepseek-v3-h200/v1",
            default_headers={"RITS_API_KEY": ""}
        )
        
        # Recreate the chain with updated LLM
        self.forward_chain = LLMChain(
            llm=self.llm,
            prompt=PromptTemplate.from_template(self.forward_prompt_template)
        )

    def forward(self, problem, comment_pool, max_collaborate_nums):
        """
        Select the next expert to consult based on problem and current state.
        
        Args:
            problem: Problem dictionary with description
            comment_pool: CommentPool containing all expert comments
            max_collaborate_nums: Maximum number of collaborations allowed
            
        Returns:
            BaseExpert: Selected expert object
        """
        try:
            all_experts = comment_pool.all_experts
            all_experts_name = [e.name for e in all_experts]
            commented_experts_name = [c.expert.name for c in comment_pool.comments]

            experts_info = '\n'.join([str(e) for e in all_experts])
            commented_experts = str(commented_experts_name)     
            remaining_experts = str(list(set(all_experts_name) - set(commented_experts_name)))
            
            # Generate expert selection
            answer = self.forward_chain.predict(
                problem_description=problem['description'], 
                experts_info=experts_info,
                commented_experts=commented_experts,
                remaining_experts=remaining_experts,
                max_collaborate_nums=max_collaborate_nums,
                remaining_collaborate_nums=max_collaborate_nums-len(commented_experts_name),
            )
            
            # Map expert names to objects
            expert_name_to_obj = { e.name: e for e in all_experts }
            
            # Find the selected expert by name matching
            selected_expert = None
            for name, expert in expert_name_to_obj.items():
                if name in answer:
                    selected_expert = expert
                    break
            
            if selected_expert is not None:
                return selected_expert
            else:
                print(f'Could not parse expert selection from: {answer}')
                print('Available experts:', list(expert_name_to_obj.keys()))
                print('Falling back to random selection')
                
                # Fallback: select from remaining experts
                remaining_expert_objs = [
                    expert for expert in all_experts 
                    if expert.name not in commented_experts_name
                ]
                
                if remaining_expert_objs:
                    return random.choice(remaining_expert_objs)
                else:
                    # If all experts have commented, select randomly
                    return random.choice(all_experts)
                    
        except Exception as e:
            print(f'Error in conductor selection: {str(e)}')
            print('Falling back to random expert selection')
            
            # Emergency fallback
            all_experts = comment_pool.all_experts
            if all_experts:
                return random.choice(all_experts)
            else:
                raise RuntimeError("No experts available for selection")