import traceback
import gradio, os
from gradio_client import Client
from dotenv import load_dotenv

load_dotenv()

client = None

def redirect(query, collection_name):
    global client

    PRIVATE_SPACE_ID = os.getenv("PRIVATE_SPACE_ID")
    PRIVATE_API_KEY = os.getenv("PRIVATE_API_KEY")

    try:
        if client is None:
            client = Client(PRIVATE_SPACE_ID, hf_token=PRIVATE_API_KEY)
    except Exception as e:
        print(f"Failed to connect to the client: {e}")
        return "Failed to connect to the client. Please try again."

    try:
        result = client.predict(
            query,
            collection_name,
            api_name="/predict"
        )
    except Exception as e:
        print(f"Failed to get prediction: {e}")
        print(f"Exception type: {type(e)}")
        print(traceback.format_exc())
        return f"Failed to get prediction ({e}), Please try again."

    return result


gradio_interface = gradio.Interface(
  fn=redirect,
  api_name="search",
  inputs=["text", "text"],
  outputs="text",
  examples=[
      ["Piso", "real_bert"],
  ],
  title="TFG Web Demo Public REST API",
  description="Used by our static web app to share the real-bert model and evaluate its performance in a real-world scenario.",
  article="© Fernando Ónega Rodrigo 2024"
)

gradio_interface.launch()