from gradio_client import Client
import os

class AussieHubClient:
    def __init__(self, space_url="Brettapps/brettapps-aussie-mcp-server-agents-mate", hf_token=None):
        self.client = Client(space_url, token=hf_token)
        self.history = []

    def chat(self, message):
        """
        Sends a message to the Aussie Hub and returns the agent response.
        Handles the internal Gradio /user -> /bot trigger sequence.
        """
        print(f"Sending to Hub: {message}")
        
        # 1. Trigger the /user endpoint to add message to history
        # returns (empty_textbox, updated_history)
        _, self.history = self.client.predict(
            user_message=message,
            history=self.history,
            api_name="/user"
        )
        
        # 2. Trigger the /bot endpoint to generate response
        # returns updated_history with agent reply
        self.history = self.client.predict(
            history=self.history,
            api_name="/bot"
        )
        
        # Extract the last message content (agent's reply)
        # Note: In Gradio 5+, history is a list of complex dicts
        last_turn = self.history[-1]
        if last_turn['role'] == 'assistant':
            return last_turn['content'][0]['text']
        return "No response received."

    def clear(self):
        """Clears the session history."""
        self.history = []
        print("Hub session cleared.")

if __name__ == "__main__":
    # Example Usage
    # Ensure HF_TOKEN is in your environment
    token = os.environ.get("HF_TOKEN")
    
    hub = AussieHubClient(hf_token=token)
    
    # Task the Hub
    response = hub.chat("G'day! Who are the core agents in this hub?")
    print(f"\nAussie Agent: {response}")
    
    # Task the Business Manager
    response = hub.chat("Check the status of Fair Dinkum Publishing.")
    print(f"\nAussie Agent: {response}")