import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load a publicly accessible model
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")

def get_alternative_xpaths(html_structure, original_xpath):
    prompt = f"""
    Given the following HTML structure and an original XPath, generate alternative XPaths that could help locate the same element if it has changed position.

    HTML Structure:
    {html_structure}

    Original XPath: {original_xpath}

    Suggested alternative XPaths:
    """
    # Generate response
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs['input_ids'], max_length=200, temperature=0.3)
    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    alternative_xpaths = response_text.strip().split('\n')
    return alternative_xpaths

def locate_element_with_self_healing(html_structure, original_xpath):
    print("Original XPath failed. Attempting to heal...")
    alternative_xpaths = get_alternative_xpaths(html_structure, original_xpath)
    
    return {
        "original_xpath": original_xpath,
        "alternative_xpaths": alternative_xpaths
    }

interface = gr.Interface(
    fn=locate_element_with_self_healing,
    inputs=[
        gr.Textbox(label="HTML Structure", placeholder="Paste the HTML structure here..."),
        gr.Textbox(label="Original XPath", placeholder="//div[@id='submit-button']")
    ],
    outputs=[
        gr.JSON(label="Healing Result")
    ],
    title="Self-Healing XPath API",
    description="This tool provides alternative XPaths if the original XPath is not found in the given HTML structure."
)

interface.launch()