from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

# Load the CodeT5 small model (lightweight)
model_name = "Salesforce/codet5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def fix_output(output):
    # Simple cleanups and common fixes
    output = output.replace("importsync_playwright", "import sync_playwright")
    output = output.replace("\n\n", "\n")
    lines = output.split("\n")
    seen = set()
    filtered = []
    for line in lines:
        if line.strip() not in seen:
            filtered.append(line)
            seen.add(line.strip())
    output = "\n".join(filtered)
    output = output.replace('.locator("//button[@type=\'submit\']").fill(', '.locator("//button[@type=\'submit\']").click(')
    return output

def convert_selenium_to_playwright(selenium_code):
    prompt = f"""
Convert the following Selenium Python code to Playwright Python code.
Return complete working Playwright code with proper syntax, including imports, browser launch, navigation, actions, and close.

Example:

Selenium:
from selenium import webdriver
from selenium.webdriver.common.by import By
driver = webdriver.Chrome()
driver.get("https://example.com")
element = driver.find_element(By.ID, "username")
element.send_keys("myusername")
driver.quit()

Playwright:
from playwright.sync_api import sync_playwright

with sync_playwright() as p:
    browser = p.chromium.launch()
    page = browser.new_page()
    page.goto("https://example.com")
    page.locator("#username").fill("myusername")
    browser.close()

Now convert:

Selenium:
{selenium_code}

Playwright:
"""
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
    outputs = model.generate(**inputs, max_length=512, num_beams=5, early_stopping=True)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return fix_output(result)

iface = gr.Interface(
    fn=convert_selenium_to_playwright,
    inputs=gr.Textbox(lines=15, placeholder="Paste your Selenium Python code here..."),
    outputs="textbox",
    title="Selenium to Playwright Code Converter",
    description="Paste your Selenium Python code and get the equivalent Playwright Python code."
)

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