import os
import sys
from huggingface_hub import InferenceClient

class ChatBot:
    def __init__(self, model_name="HelpingAI/HELVETE-3B"):
        """Initialize the chatbot with the specified model."""
        # Get API token from environment variable
        self.api_key = os.environ.get("HF_TOKEN")
        if not self.api_key:
            print("Error: HF_TOKEN environment variable not set")
            print("Please set your Hugging Face token with: export HF_TOKEN=your_token_here")
            sys.exit(1)
        
        self.client = InferenceClient(api_key=self.api_key)
        self.model = model_name
        self.conversation_history = []
    
    def add_message(self, role, content):
        """Add a message to the conversation history."""
        self.conversation_history.append({"role": role, "content": content})
    
    def get_response(self, user_message):
        """Get a response from the AI model."""
        # Add user message to history
        self.add_message("user", user_message)
        
        try:
            # Get completion from the model
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=self.conversation_history,
            )
            
            # Extract the assistant's response
            assistant_response = completion.choices[0].message.content
            
            # Add assistant response to history
            self.add_message("assistant", assistant_response)
            
            return assistant_response
            
        except Exception as e:
            return f"Error: {str(e)}"
    
    def chat_loop(self):
        """Run an interactive chat loop."""
        print("=" * 50)
        print("🤖 HELVETE-3B ChatBot")
        print("Type 'quit', 'exit', or 'bye' to end the conversation")
        print("=" * 50)
        
        while True:
            try:
                user_input = input("\nYou: ").strip()
                
                if user_input.lower() in ['quit', 'exit', 'bye']:
                    print("Goodbye! 👋")
                    break
                
                if not user_input:
                    continue
                
                print("AI: ", end="", flush=True)
                response = self.get_response(user_input)
                print(response)
                
            except KeyboardInterrupt:
                print("\n\nConversation interrupted. Goodbye! 👋")
                break
            except Exception as e:
                print(f"\nAn error occurred: {e}")

def main():
    """Main function to run the chatbot."""
    # Create and run the chatbot
    bot = ChatBot()
    bot.chat_loop()

if __name__ == "__main__":
    main()