import gradio as gr
import requests

# Set your Hugging Face API key here
API_KEY = "your_huggingface_api_key"

def unified_function(operation, input_text):
    if operation == "Clean Diff":
        added_lines = []
        for line in input_text.split('\n'):
            if line.startswith('+') and not line.startswith('+++'):
                added_line = line[1:]  # Remove the '+' sign
                if sum(len(l) + 1 for l in added_lines) + len(added_line) <= 2000:
                    added_lines.append(added_line)
                else:
                    return "too long, try again"
                    break
        return '\n'.join(added_lines)

    elif operation == "Classify":
        API_URL = "https://api-inference.huggingface.co/models/davidgaofc/TechDebtClassifier"
        headers = {"Authorization": f"Bearer {API_KEY}"}
        data = {"inputs": input_text}

        response = requests.post(API_URL, headers=headers, json=data)
        result = response.json()
        return result

    elif operation == "Generate Label":
        API_URL = "https://api-inference.huggingface.co/models/davidgaofc/TechDebtLabeler"
        headers = {"Authorization": f"Bearer {API_KEY}"}
        data = {"inputs": input_text}

        response = requests.post(API_URL, headers=headers, json=data)
        result = response.json()
        return result

def huggingface_login(api_key):
    global API_KEY
    API_KEY = api_key
    return "success!"
# Create the Gradio interface
interface = gr.Interface(
    fn=unified_function,
    inputs=[
        gr.Dropdown(["Clean Diff", "Classify", "Generate Label"], label="Select Operation"),
        gr.Textbox(label="Input Text")
    ],
    outputs="text",
    title="Unified Interface for Multiple Functions",
    description="Select an operation from the dropdown and input text to see the result."
)

huggingface_interface = gr.Interface(
    fn=huggingface_login,
    inputs=gr.Textbox(lines=1, label="API Key"),
    outputs = "text"
)

tabbed_interface = gr.TabbedInterface([huggingface_interface, interface], ["Login", "Main"])

if __name__ == "__main__":
    tabbed_interface.launch()