import gradio as gr
import gspread
from oauth2client.service_account import ServiceAccountCredentials
import json
import os
import requests

# Load Google Sheets credentials from secrets
creds_json = os.getenv("GOOGLE_SHEETS_KEY_JSON")
creds_dict = json.loads(creds_json)

# Google Sheets setup
scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
creds = ServiceAccountCredentials.from_json_keyfile_dict(creds_dict, scope)
client = gspread.authorize(creds)
sheet = client.open("QF - Quant Request").sheet1  # Use the correct sheet name

def validate_repo(link):
    """Check if the repository exists on Hugging Face."""
    # Extract repo_id from the link (format: {username}/{model_name} or {dataset_name})
    repo_id = link.split("https://huggingface.co/")[-1].strip("/")
    url = f"https://huggingface.co/api/models/{repo_id}"  # Adjust the URL for datasets if needed
    
    response = requests.head(url)
    return response.status_code == 200

def check_quant_exists(link):
    """Check if the quantization request already exists in the Google Sheet."""
    extracted_text = link.split("https://huggingface.co/")[-1]
    existing_texts = sheet.col_values(1)  # Assuming the extracted text is in the first column
    return extracted_text in existing_texts

def check_quant_factory_quant_exists(link):
    """Check if a GGUF quantized model exists in the QuantFactory repository."""
    # Extract the model name from the original link
    model_name = link.split('/')[-1]  # Get the last item from the split list, which is the model name

    # Create the QuantFactory GGUF quant link
    quant_factory_link = f"https://huggingface.co/QuantFactory/{model_name}-GGUF"

    # Check if the QuantFactory GGUF repo exists
    response = requests.head(quant_factory_link)

    if response.status_code == 200:
        return True, quant_factory_link
    else:
        return False, quant_factory_link

def submit_link(link):
    # Normalize the input link
    if not link.startswith("https://huggingface.co"):
        link = f"https://huggingface.co/{link}"
    
    # Validate the repo
    if not validate_repo(link):
        return "Invalid model or repository link. Please provide a valid Hugging Face link."

    # Check if the GGUF quantized model already exists in QuantFactory
    quant_exists, quant_link = check_quant_factory_quant_exists(link)
    if quant_exists:
        model_name = link.split('/')[-1]  # Extract the model name
        # Return a formatted Markdown string to create a clickable hyperlink
        return f"Quant already exists at [QuantFactory/{model_name}-GGUF]({quant_link})."

    # Check if the quant request already exists in the Google Sheet
    if check_quant_exists(link):
        return "Quant requests have already been made for this model."

    # Extract text after "huggingface.co/"
    extracted_text = link.split("https://huggingface.co/")[-1]
    
    # Append the row with the extracted text and default status "Pending"
    row_index = len(sheet.get_all_values()) + 1
    sheet.append_row(["", ""])  # Append an empty row first to ensure correct row index
    
    # Set the hyperlink formula in the first column
    sheet.update_cell(row_index, 1, f'=HYPERLINK("{link}", "{extracted_text}")')
    
    # Copy content from cell B2 to the new row in column B
    b2_value = sheet.cell(2, 2).value
    sheet.update_cell(row_index, 2, b2_value)

    return "Request submitted successfully."

# Gradio Interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):  # Left column
            gr.Markdown("## QuantFactory - Model Request")
            gr.Markdown("*NOTE: Due to limited resources, currently we can only quantize models with upto 20B parameters.*")
            link_input = gr.Textbox(label="Hugging Face Link")
            submit_button = gr.Button("Submit")
            result = gr.Markdown()  # Use Markdown to render the output as a clickable hyperlink
            submit_button.click(fn=submit_link, inputs=link_input, outputs=result)
        with gr.Column(scale=1):  # Right column
            gr.Image("image.png")  # Display the image on the right

demo.launch()