# coding: utf-8

class Agent:
    def __init__(self, model_caller) -> None:
        self.model_caller = model_caller
        self.memory = [{"role": "system", "content": "You are a social science expert."}]

    def generate_messages(self, prompt):
        self.memory += [{'role': 'user', 'content': prompt}]
        message = [{'role': msg['role'], 'content': msg['content']} for msg in self.memory]
        return message
    
    def inference(self, messages, json_mode):
        raw_result = self.model_caller.generate(messages, json_mode).strip()
        self.memory += [{'role': 'assistant', 'content': raw_result}]
        return raw_result

class Chair(Agent):
    def __init__(self, model_caller) -> None:
        super().__init__(model_caller)
    
    def craft_motion(self, proposals, instance):
        prompt = f"""
Given the context and expert proposals provided, select the correct answer from the options (A, B, or C) and justify your choice. 
Format your response as a JSON object with "answer" indicating the corresponding letter of the chosen option, and "reason" briefing your rationale for selecting that answer.

Expert Proposals:
{proposals}

Context: {instance['context']}
Question: {instance['question']}
A) {instance['ans0']}
B) {instance['ans1']}
C) {instance['ans2']}
Answer JSON:"""
        messages = self.generate_messages(prompt)
        motion = self.inference(messages, True)
        return motion

class Member(Agent):    
    def __init__(self, bias_type, model_caller) -> None:
        super().__init__(model_caller)
        self.bias_type = bias_type

    def propose(self, instance):
        prompt = f"""
Generate concise guidelines to prevent {self.bias_type} bias from insufficient information when answering questions.
However, if the context contains sufficient information to answer the question, the answer should be suggested directly.

Context: {instance['context']} 
Question: {instance['context']}

Your propose:"""
        messages = self.generate_messages(prompt)
        proposal = self.inference(messages, False)
        return proposal
    
    def vote(self, motion, instance):
        prompt = f"""
Pass the motion if the answer is successfully prevented from {self.bias_type} bias. 
Otherwise, Veto the motion. You can abstain if the current motion does not involve {self.bias_type} bias. 
Response your decision in the JSON format: {{"decision": "<Pass/Veto/Abstain>"}}

Context: {instance['context']}
Question: {instance['question']}
A) {instance['ans0']}
B) {instance['ans1']}
C) {instance['ans2']}

Motion: {motion}

Your decision:"""
        messages = self.generate_messages(prompt)
        vote = self.inference(messages, True)
        return vote

class Committee:
    def __init__(self, committee_dict) -> None:
        self.chair = Chair(committee_dict['chair']['model_caller'])
        self.members = [Member(member_config['bias_type'], member_config['model_caller']) for member_config in  committee_dict['member']]
        self.current_motion = None

    def deliberate(self, instance):
        # Propose
        proposals = [member.propose(instance) for member in self.members]

        # Craft Motion
        motion = self.chair.craft_motion(proposals, instance)

        # Vote
        vote = [member.vote(motion, instance) for member in self.members]

        return vote