import os
import gradio as gr
import chromadb
import json
from typing import List, Dict
from huggingface_hub import snapshot_download

# Initialize ChromaDB client
client = chromadb.PersistentClient(path=str(snapshot_download("mrfakename/librivox-db", token=os.getenv("HF_TOKEN"))))
collection = client.get_or_create_collection("librivox_catalog")

# Custom CSS for modern UI
custom_css = """
.book-search-container {
    max-width: 1200px;
    margin: 0 auto;
    padding: 20px;
}
.search-results {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
    gap: 20px;
    margin-top: 20px;
}
@media (min-width: 1024px) {
    .search-results {
        grid-template-columns: repeat(4, 1fr);
    }
}
@media (max-width: 1023px) {
    .search-results {
        grid-template-columns: repeat(3, 1fr);
    }
}
@media (max-width: 768px) {
    .search-results {
        grid-template-columns: repeat(2, 1fr);
    }
}
@media (max-width: 480px) {
    .search-results {
        grid-template-columns: 1fr;
    }
}
.book-card {
    border-radius: 10px;
    padding: 15px;
    background: var(--background-fill-primary);
    border: 1px solid var(--border-color-primary);
    transition: all 0.2s;
    text-decoration: none;
    color: inherit;
    display: block;
}
.book-card:hover {
    transform: translateY(-5px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.1);
}
.book-cover {
    width: 100%;
    height: 200px;
    object-fit: cover;
    border-radius: 5px;
}
.book-title {
    font-size: 1.1em;
    font-weight: bold;
    margin: 10px 0 5px;
    line-height: 1.3;
}
.book-meta {
    font-size: 0.9em;
    color: var(--body-text-color-secondary);
    display: flex;
    flex-wrap: wrap;
    gap: 8px;
}
.meta-item {
    display: inline-flex;
    align-items: center;
    gap: 4px;
}
"""

def format_duration(seconds: int) -> str:
    hours = seconds // 3600
    minutes = (seconds % 3600) // 60
    return f"{hours}h {minutes}m"

def search_books(query: str, n_results: int = 10) -> str:
    if not query.strip():
        return "Please enter a search query"
        
    results = collection.query(
        query_texts=[query],
        n_results=n_results
    )
    
    html_results = []
    for idx, (doc, metadata) in enumerate(zip(results['documents'][0], results['metadatas'][0])):
        metadata['authors'] = json.loads(metadata['authors'])
        metadata['genres'] = json.loads(metadata['genres'])
        
        authors = ", ".join([author['last_name'] + ", " + author['first_name'] for author in metadata['authors']])
        genres = ", ".join([genre['name'] for genre in metadata['genres']])
        duration = format_duration(int(metadata['totaltimesecs']))
        
        card_html = f"""
        <a href="{metadata['url_librivox']}" target="_blank" class="book-card">
            <img class="book-cover" src="{metadata['coverart_jpg']}" alt="Cover for {metadata['title']}">
            <div class="book-title">{metadata['title']}</div>
            <div class="book-meta">
                <span class="meta-item">📚 {authors}</span>
                <span class="meta-item">⏱️ {duration}</span>
                <span class="meta-item">🎭 {genres}</span>
                <span class="meta-item">🌐 {metadata['language']}</span>
            </div>
        </a>
        """
        html_results.append(card_html)
    
    return f"""
    <div class="book-search-container">
        <div class="search-results">
            {''.join(html_results)}
        </div>
    </div>
    """

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("# LibriVox Semantic Search")
    
    with gr.Group():
        query = gr.Textbox(
            label="Search Query",
            placeholder="Enter keywords to search audiobooks...",
        )
        n_results = gr.Slider(
            minimum=5,
            maximum=50,
            value=10,
            step=5,
            label="Number of Results",
        )
    
    search_button = gr.Button("Search", variant="primary")
    results = gr.HTML()
    
    search_button.click(
        fn=search_books,
        inputs=[query, n_results],
        outputs=results
    )
    
    # Also trigger search on enter key
    query.submit(
        fn=search_books,
        inputs=[query, n_results],
        outputs=results
    )

if __name__ == "__main__":
    demo.launch()