# app.py

import os
import time
import gradio as gr
import importlib.util

from huggingface_hub import hf_hub_download

# ----------------------------------------------------------------------
# Helper to read secrets from the HF Space environment
# ----------------------------------------------------------------------
def _secret(key: str, fallback: str = None) -> str:
    val = os.getenv(key)
    if val is not None: return val
    if fallback is not None: return fallback
    raise RuntimeError(f"Secret '{key}' not found. Please add it to your Space secrets.")

# ----------------------------------------------------------------------
# 1. Configuration & Constants
# ----------------------------------------------------------------------
# The private repo containing the vector DB and the logic script
REPO_ID = _secret("REPO_ID") 
# Files to download from the repo
FILES_TO_DOWNLOAD = ["index.faiss", "index.pkl", "agent_logic.py"]
# A local directory to store all downloaded assets
LOCAL_DOWNLOAD_DIR = "downloaded_assets"
EMBEDDING_MODEL_NAME = "google/embeddinggemma-300m" 

# ----------------------------------------------------------------------
# 2. Bootstrap Phase: Download assets and initialize the engine
#    (This code runs only once when the Space starts up)
# ----------------------------------------------------------------------
print("--- [UI App] Starting bootstrap process ---")
os.makedirs(LOCAL_DOWNLOAD_DIR, exist_ok=True)
hf_token = _secret("HF_TOKEN") # A read-access token is required for private repos

for filename in FILES_TO_DOWNLOAD:
    print(f"--- [UI App] Downloading '{filename}'... ---")
    try:
        hf_hub_download(
            repo_id=REPO_ID, filename=filename, repo_type="dataset",
            local_dir=LOCAL_DOWNLOAD_DIR, token=hf_token,
        )
    except Exception as e:
        raise RuntimeError(f"Failed to download '{filename}'. Check repo/file names and HF_TOKEN. Error: {e}")

# Dynamically import the RAG_Engine class from the downloaded script
logic_script_path = os.path.join(LOCAL_DOWNLOAD_DIR, "agent_logic.py")
spec = importlib.util.spec_from_file_location("agent_logic", logic_script_path)
agent_logic_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(agent_logic_module)
print("--- [UI App] Agent logic module imported successfully. ---")

# Instantiate the engine. This single line triggers all the complex setup
# defined in the private_logic.py file.
engine = agent_logic_module.RAG_Engine(
    local_download_dir=LOCAL_DOWNLOAD_DIR,
    embedding_model_name=EMBEDDING_MODEL_NAME
)
print("--- [UI App] Bootstrap complete. Gradio UI is starting. ---")

# ----------------------------------------------------------------------
# 3. Core Gradio Chat Logic (Now a simple wrapper)
# ----------------------------------------------------------------------
def respond(message: str, history: list[dict[str, str]]):
    """
    This function is called by Gradio for each user message.
    It passes the inputs to the RAG engine and streams the output.
    """
    final_response = engine.get_response(message, history)
    
    # Stream the response back to the UI for a "typing" effect
    response = ""
    for char in final_response:
        response += char
        time.sleep(0.01)
        yield response

# ----------------------------------------------------------------------
# 4. UI Layout and Launch
# ----------------------------------------------------------------------
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    title="PRECISE RAG Agent",
    description="Silakan bertanya tentang PRECISE.",
    examples=[
        ["Apa rumus untuk menghitung PVR?"],
        ["Apa tujuan pengadaan PRECISE?"],
    ],
    cache_examples=False,
    theme=gr.themes.Soft(),
)

with gr.Blocks() as demo:
    chatbot.render()

if __name__ == "__main__":
    allowed_user = _secret("CHAT_USER")
    allowed_pass = _secret("CHAT_PASS")
    demo.launch(
        auth=(allowed_user, allowed_pass),
        server_name="0.0.0.0",
        ssr_mode=False,
        server_port=7860
    )