# app.py

import os
import time
import gradio as gr
import importlib.util
import markdown
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. Custom CSS for better styling
# ----------------------------------------------------------------------
custom_css = """
/* Main container styling */
.contain {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    min-height: 100vh;
    padding: 20px;
}

/* Header styling */
.header {
    text-align: center;
    margin-bottom: 30px;
}

.header h1 {
    color: #2c3e50;
    font-weight: 700;
    margin-bottom: 10px;
    font-size: 2.5rem;
}

.header p {
    color: #7f8c8d;
    font-size: 1.2rem;
}

/* Chat container styling */
.gr-chat-message {
    padding: 16px 20px;
    border-radius: 18px;
    margin: 10px 0;
    line-height: 1.5;
}

.gr-chat-message-user {
    background: linear-gradient(135deg, #3498db 0%, #2980b9 100%);
    color: white;
    margin-left: 20%;
}

.gr-chat-message-bot {
    background: linear-gradient(135deg, #ecf0f1 0%, #ffffff 100%);
    color: #2c3e50;
    margin-right: 20%;
    border: 1px solid #e0e0e0;
}

/* Chat input styling */
.gr-text-input textarea {
    border-radius: 12px;
    padding: 15px;
    font-size: 16px;
    border: 2px solid #bdc3c7;
}

.gr-text-input textarea:focus {
    border-color: #3498db;
    box-shadow: 0 0 0 3px rgba(52, 152, 219, 0.2);
}

/* Button styling */
.gr-button {
    background: linear-gradient(135deg, #2ecc71 0%, #27ae60 100%);
    color: white;
    border: none;
    border-radius: 12px;
    padding: 12px 24px;
    font-weight: 600;
    transition: all 0.3s ease;
}

.gr-button:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}

/* Example styling */
.gr-examples {
    border-radius: 12px;
    background: white;
    padding: 15px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}

.gr-examples button {
    background: #f8f9fa;
    border: 1px solid #e9ecef;
    border-radius: 8px;
    padding: 10px 15px;
    margin: 5px;
    transition: all 0.2s ease;
}

.gr-examples button:hover {
    background: #e9ecef;
    transform: translateY(-1px);
}

/* Formula styling */
.formula {
    background-color: #f8f9fa;
    border-left: 4px solid #3498db;
    padding: 15px;
    margin: 15px 0;
    border-radius: 0 8px 8px 0;
    overflow-x: auto;
}

.formula p {
    margin: 0;
    font-family: monospace;
    color: #2c3e50;
}

/* Responsive adjustments */
@media (max-width: 768px) {
    .gr-chat-message-user, .gr-chat-message-bot {
        margin-left: 10%;
        margin-right: 10%;
    }
    
    .header h1 {
        font-size: 2rem;
    }
}
"""

# ----------------------------------------------------------------------
# 5. UI Layout and Launch
# ----------------------------------------------------------------------
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    # Header section
    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML("""
            <div class="header">
                <h1>🤖 PRECISE RAG Agent</h1>
                <p>Silakan bertanya tentang PRECISE dan sistem promosinya</p>
            </div>
            """)
    
    # Chat interface
    chatbot = gr.ChatInterface(
        respond,
        type="messages",
        examples=[
            ["Apa rumus untuk menghitung PVR?"],
            ["Apa tujuan pengadaan PRECISE?"],
            ["Jelaskan tentang promo 'beli 4 gratis 2'"],
        ],
        cache_examples=False,
    )
    
    # Additional information section
    with gr.Accordion("📊 Informasi Formula PVR", open=False):
        gr.Markdown("""
        **PVR untuk promo "beli 4 gratis 2" dihitung dengan rumus:**
        
        ```math
        \\text{PVR} = \\frac{\\text{added\\_value}}{\\text{purchase\\_value}} 
                    = \\frac{\\text{target\\_value} \\times \\text{avg\\_price\\_AVP}}
                           {\\text{target\\_value} \\times \\text{avg\\_price\\_AVP} 
                            + \\text{condition\\_value} \\times \\text{avg\\_price\\_REP}}
        ```
        
        Keterangan:
        - **added_value**: Nilai tambah dari promo
        - **purchase_value**: Nilai pembelian keseluruhan
        - **target_value**: Jumlah produk target dalam promo
        - **avg_price_AVP**: Harga rata-rata produk AVP
        - **condition_value**: Jumlah produk kondisi dalam promo
        - **avg_price_REP**: Harga rata-rata produk REP
        """)

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
    )