import argparse
import os
import requests
from dotenv import load_dotenv

load_dotenv()  # Load environment variables from .env

def generate_text(prompt):
    # Use the correct model URL for EleutherAI's GPT-J 6B
    url = "https://api-inference.huggingface.co/models/EleutherAI/gpt-j-6b"
    headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY')}"}
    
    # Send the request with the prompt as input
    response = requests.post(url, headers=headers, json={"inputs": prompt})
    
    # Check the response status and return the generated text if successful
    if response.status_code == 200:
        return response.json()[0]["generated_text"]  # Adjust for API's JSON structure
    else:
        return f"Error: {response.status_code} - {response.text}"

def cli_interface():
    parser = argparse.ArgumentParser(description="Command-line interaction with the deployed model.")
    parser.add_argument("--task", type=str, help="The prompt or command to generate text for")
    args = parser.parse_args()

    # Use the provided prompt or default to "Tell me a joke"
    task = args.task if args.task else "Tell me a joke"
    result = generate_text(task)
    print(result)

if __name__ == "__main__":
    cli_interface()