import os
import json
import gradio as gr
import lancedb
import torch
# Limit PyTorch CPU threads to 1 to prevent system freezing
torch.set_num_threads(1)
from huggingface_hub import snapshot_download
from sentence_transformers import SentenceTransformer

# 1. Setup paths and download dataset from Hugging Face Hub
DB_REPO = "anurag-chand/vivekananda-scriptures-lancedb"
LOCAL_DB_DIR = "./scriptures_lancedb"

print(f"Checking for scriptures database locally at {LOCAL_DB_DIR}...")
if not os.path.exists(LOCAL_DB_DIR) or not os.listdir(LOCAL_DB_DIR):
    print(f"Database not found. Downloading {DB_REPO} from Hugging Face Hub...")
    snapshot_download(
        repo_id=DB_REPO,
        repo_type="dataset",
        local_dir=LOCAL_DB_DIR
    )
    print("Download complete!")

# Connect to LanceDB
db = lancedb.connect(LOCAL_DB_DIR)
table = db.open_table("scriptures")
print(f"Connected to scriptures table! Total records: {len(table):,}")

# 2. Load model
print("Loading Krutrim Vyakyarth model...")
model = SentenceTransformer("krutrim-ai-labs/Vyakyarth", device="cpu")
print("Model loaded successfully.")

# Custom CSS for gorgeous aesthetics (dark mode, glassmorphism, responsive grid)
custom_css = """
.header-box {
    text-align: center;
    padding: 30px 20px;
    background: linear-gradient(135deg, rgba(30, 41, 59, 0.5), rgba(15, 23, 42, 0.8));
    border: 1px solid rgba(255, 255, 255, 0.05);
    border-radius: 16px;
    margin-bottom: 30px;
    box-shadow: 0 4px 30px rgba(0, 0, 0, 0.3);
}
.header-box h1 {
    font-size: 2.5em;
    font-weight: 800;
    background: linear-gradient(to right, #ffd700, #ff8c00);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 10px;
}
.header-box p {
    color: #94a3b8;
    font-size: 1.1em;
}
.search-btn {
    background: linear-gradient(135deg, #ff8c00, #d35400) !important;
    color: white !important;
    font-weight: bold !important;
    border: none !important;
    border-radius: 8px !important;
    transition: all 0.3s ease !important;
}
.search-btn:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 4px 15px rgba(211, 84, 0, 0.4) !important;
}
.result-card {
    background: rgba(30, 41, 59, 0.4);
    border: 1px solid rgba(255, 255, 255, 0.06);
    border-radius: 14px;
    padding: 22px;
    margin-bottom: 20px;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.15);
    transition: all 0.3s ease;
}
.result-card:hover {
    transform: translateY(-2px);
    border-color: rgba(255, 140, 0, 0.3);
    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
}
.card-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    border-bottom: 1px solid rgba(255, 255, 255, 0.06);
    padding-bottom: 10px;
    margin-bottom: 15px;
    flex-wrap: wrap;
    gap: 10px;
}
.match-badge {
    background: linear-gradient(135deg, #ff8c00, #e67e22);
    color: white;
    padding: 4px 12px;
    border-radius: 20px;
    font-size: 0.85em;
    font-weight: bold;
}
.meta-tags {
    display: flex;
    gap: 8px;
    flex-wrap: wrap;
}
.meta-tag {
    background: rgba(255, 255, 255, 0.05);
    border: 1px solid rgba(255, 255, 255, 0.08);
    color: #cbd5e1;
    padding: 3px 10px;
    border-radius: 6px;
    font-size: 0.85em;
}
.shloka-section {
    background: rgba(255, 215, 0, 0.03);
    border-left: 4px solid #ffd700;
    padding: 12px 16px;
    margin-bottom: 15px;
    border-radius: 0 8px 8px 0;
}
.shloka-text {
    font-size: 1.25em;
    color: #f1c40f;
    line-height: 1.6;
    margin: 0;
    font-weight: bold;
}
.translation-section {
    background: rgba(46, 204, 113, 0.03);
    border-left: 4px solid #2ecc71;
    padding: 12px 16px;
    margin-bottom: 15px;
    border-radius: 0 8px 8px 0;
}
.translation-text {
    font-size: 1.05em;
    color: #2ecc71;
    line-height: 1.5;
    margin: 0;
    font-style: italic;
}
.commentary-section {
    padding-left: 4px;
}
.commentary-header {
    font-size: 0.95em;
    font-weight: bold;
    color: #94a3b8;
    margin-bottom: 6px;
    text-transform: uppercase;
    letter-spacing: 0.05em;
}
.commentary-text {
    font-size: 1.05em;
    color: #cbd5e1;
    line-height: 1.6;
    margin: 0;
}
"""

def semantic_search(query: str, limit: int = 5, min_similarity: float = 0.5) -> str:
    if not query.strip():
        return "<div style='text-align: center; color: #94a3b8; padding: 20px;'>Please enter a search query above.</div>"
        
    try:
        # Encode query
        query_vector = model.encode(query).tolist()
        
        # Search LanceDB table with Cosine Similarity
        results = table.search(query_vector).metric("cosine").limit(limit).to_list()
        
        if not results:
            return "<div style='text-align: center; color: #ff6b6b; padding: 20px;'>No results found. Try adjusting your query.</div>"
            
        html_output = ""
        for idx, res in enumerate(results):
            # Parse metadata
            meta = {}
            if "metadata" in res and res["metadata"]:
                try:
                    meta = json.loads(res["metadata"])
                except Exception:
                    pass
                    
            # Compute similarity from Cosine distance
            distance = res.get("_distance", 1.0)
            similarity = max(0.0, 1.0 - distance)
            
            if similarity < min_similarity:
                continue
                
            source = res.get("source_file") or meta.get("source_file") or "Unknown"
            book = meta.get("book_title") or meta.get("title") or os.path.splitext(source)[0]
            book = str(book).replace("_", " ").strip().title()
            
            chapter = meta.get("chapter_title") or meta.get("section") or meta.get("chapter_label") or ""
            chapter = str(chapter).replace("_", " ").strip().title()
            
            shloka = meta.get("shloka") or meta.get("sutra") or ""
            translation = meta.get("translation") or ""
            commentary_author = meta.get("commentary_author") or ""
            raw_text = res.get("text", "")
            
            # Format shloka block
            shloka_html = ""
            if shloka:
                shloka_formatted = shloka.strip().replace("\n", "<br>")
                shloka_html = f"""
                <div class="shloka-section">
                    <p class="shloka-text">{shloka_formatted}</p>
                </div>
                """
                
            # Format translation block
            translation_html = ""
            if translation:
                translation_formatted = translation.strip().replace("\n", "<br>")
                translation_html = f"""
                <div class="translation-section">
                    <p class="translation-text">{translation_formatted}</p>
                </div>
                """
                
            # Format raw text chunk exactly as stored in database
            text_formatted = raw_text.strip().replace("\n", "<br>")
            comm_title = "Full Scripture & Commentary Chunk"
            
            html_output += f"""
            <div class="result-card">
                <div class="card-header">
                    <span class="match-badge">Similarity: {similarity*100:.1f}%</span>
                    <div class="meta-tags">
                        <span class="meta-tag">📖 {book}</span>
                        {f'<span class="meta-tag">🔖 {chapter}</span>' if chapter else ''}
                        <span class="meta-tag">📄 {source}</span>
                    </div>
                </div>
                {shloka_html}
                {translation_html}
                <div class="commentary-section">
                    <div class="commentary-header">{comm_title}</div>
                    <p class="commentary-text">{text_formatted}</p>
                </div>
            </div>
            """
            
        if not html_output:
            return f"<div style='text-align: center; color: #ff6b6b; padding: 20px;'>No results matched the similarity threshold of {min_similarity*100:.0f}%. Try lowering it.</div>"
            
        return html_output
        
    except Exception as e:
        return f"<div style='text-align: center; color: #ff6b6b; padding: 20px;'>Error during search: {e}</div>"

# 3. Create Gradio Interface Block with stunning native dark theme
with gr.Blocks(theme=gr.themes.Default(primary_hue="amber", secondary_hue="orange", neutral_hue="slate", dark_mode=True), css=custom_css, title="Vivekananda Scriptures Semantic Search") as demo:
    # Header block
    gr.HTML("""
    <div class="header-box">
        <h1>Swami Vivekananda Scriptures Semantic Search</h1>
        <p>Search over 686,000+ chunks of Sanskrit scriptures, translations, and commentaries (Patanjali Yoga Sutras, Gaudapada Karika, Vivekananda lectures, and more) using high-precision SOTA semantic vector matching.</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=4):
            query_input = gr.Textbox(
                label="Search Query", 
                placeholder="Type your search here (e.g., liberation from the cycle of birth and death, control of mind, nature of Brahman)...",
                lines=1
            )
        with gr.Column(scale=1):
            search_button = gr.Button("Search", elem_classes=["search-btn"])
            
    with gr.Row():
        with gr.Column(scale=1):
            limit_slider = gr.Slider(
                label="Number of Results", 
                minimum=1, 
                maximum=20, 
                value=5, 
                step=1
            )
        with gr.Column(scale=1):
            similarity_slider = gr.Slider(
                label="Minimum Similarity Threshold", 
                minimum=0.0, 
                maximum=1.0, 
                value=0.35, 
                step=0.05
            )
            
    # Outputs block
    results_output = gr.HTML(label="Search Results")
    
    # Event binds
    search_button.click(
        fn=semantic_search,
        inputs=[query_input, limit_slider, similarity_slider],
        outputs=results_output
    )
    query_input.submit(
        fn=semantic_search,
        inputs=[query_input, limit_slider, similarity_slider],
        outputs=results_output
    )
    
if __name__ == "__main__":
    demo.launch()