"""
Blueprint Magazine Search - Gradio Interface
Search through indexed articles using semantic and keyword search.
"""

import os
import hashlib
import requests
import gradio as gr
from pathlib import Path
from datetime import datetime, timedelta
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
import bm25s
import joblib
import Stemmer
import json

# Configuration
INDEX_URL = os.environ.get('INDEX_URL', '')  # URL to encrypted index file
INDEX_ENCRYPTION_KEY = os.environ.get('INDEX_ENCRYPTION_KEY', '')
CACHE_DIR = Path('cache')
CHROMA_DB_PATH = CACHE_DIR / 'chroma_db'
BM25_INDEX_PATH = CACHE_DIR / 'bm25_index.pkl'
BM25_METADATA_PATH = CACHE_DIR / 'bm25_metadata.json'
ENCRYPTED_INDEX_PATH = CACHE_DIR / 'search_index.enc'
LAST_UPDATE_FILE = CACHE_DIR / 'last_update.txt'

# Create cache directory
CACHE_DIR.mkdir(exist_ok=True)


def get_encryption_key():
    """Get or derive 32-byte encryption key from environment."""
    key_str = INDEX_ENCRYPTION_KEY
    if not key_str:
        raise ValueError("INDEX_ENCRYPTION_KEY environment variable not set")
    
    # Derive 32-byte key using SHA-256
    return hashlib.sha256(key_str.encode()).digest()


def decrypt_directory(input_path, output_dir, key):
    """Decrypt file and extract directory."""
    import tarfile
    import io
    
    if len(key) != 32:
        raise ValueError("Key must be 32 bytes")
    
    # Read encrypted file
    with open(input_path, 'rb') as f:
        data = f.read()
    
    # Extract nonce and ciphertext
    nonce = data[:12]
    ciphertext = data[12:]
    
    # Decrypt (validates tag automatically)
    cipher = ChaCha20Poly1305(key)
    plaintext = cipher.decrypt(nonce, ciphertext, None)
    
    # Extract tar archive
    tar_buffer = io.BytesIO(plaintext)
    with tarfile.open(fileobj=tar_buffer, mode='r:gz') as tar:
        tar.extractall(path=output_dir.parent)


def should_update_index():
    """Check if index needs to be updated (24hr cache)."""
    if not LAST_UPDATE_FILE.exists():
        return True
    
    with open(LAST_UPDATE_FILE, 'r') as f:
        last_update_str = f.read().strip()
    
    try:
        last_update = datetime.fromisoformat(last_update_str)
        if datetime.now() - last_update > timedelta(hours=24):
            return True
    except:
        return True
    
    return False


def download_and_decrypt_index():
    """Download encrypted index and decrypt it."""
    print("📥 Downloading encrypted index...")
    
    try:
        # Download encrypted index
        response = requests.get(INDEX_URL, timeout=60)
        response.raise_for_status()
        
        # Save encrypted file
        with open(ENCRYPTED_INDEX_PATH, 'wb') as f:
            f.write(response.content)
        
        print(f"✓ Downloaded {len(response.content) / (1024*1024):.2f} MB")
        
        # Decrypt
        print("🔓 Decrypting index...")
        encryption_key = get_encryption_key()
        
        # Remove old database if exists
        if CHROMA_DB_PATH.exists():
            import shutil
            shutil.rmtree(CHROMA_DB_PATH)
        
        decrypt_directory(ENCRYPTED_INDEX_PATH, CHROMA_DB_PATH, encryption_key)
        print("✓ Index decrypted successfully")
        print(os.listdir(CHROMA_DB_PATH))
        
        # Update last update timestamp
        with open(LAST_UPDATE_FILE, 'w') as f:
            f.write(datetime.now().isoformat())
        
        return True
        
    except Exception as e:
        print(f"✗ Error downloading/decrypting index: {e}")
        return False


def initialize_indexes():
    """Initialize ChromaDB and BM25 indexes."""
    
    # Initialize semantic embeddings for ChromaDB
    semantic_ef = SentenceTransformerEmbeddingFunction(
        model_name="Qwen/Qwen3-Embedding-0.6B",
        device="cpu",
        normalize_embeddings=False
    )
    
    # Create persistent client
    client = chromadb.PersistentClient(path=str(CHROMA_DB_PATH))
    
    # Get collection
    semantic_collection = client.get_collection(
        name="posts_semantic",
        embedding_function=semantic_ef
    )
    
    # Load BM25 index and metadata
    bm25_retriever = joblib.load(BM25_INDEX_PATH)
    with open(BM25_METADATA_PATH, 'r') as f:
        bm25_metadata = json.load(f)
    
    return client, semantic_collection, bm25_retriever, bm25_metadata


def search_articles(query, search_type="hybrid", n_results=10):
    """
    Search articles using BM25, semantic, or hybrid search.
    
    Args:
        query: Search query string
        search_type: "keywords", "semantic", or "hybrid"
        n_results: Number of results to return
    
    Returns:
        List of dicts with article metadata
    """
    
    if not query.strip():
        return []
    
    try:
        # Check if index needs updating
        if should_update_index():
            success = download_and_decrypt_index()
            if not success and not CHROMA_DB_PATH.exists():
                return [{
                    'title': 'Error',
                    'excerpt': 'Failed to load search index. Please try again later.',
                    'url': '',
                    'tags': '',
                    'published_at': '',
                    'feature_image': '',
                    'score': 0.0
                }]
        
        # Initialize indexes
        client, semantic_collection, bm25_retriever, bm25_metadata = initialize_indexes()
        
        results_dict = {}  # Use dict to merge results by ID
        
        # BM25 keyword search
        if search_type in ["keywords", "hybrid"]:
            stemmer = Stemmer.Stemmer("english")
            query_tokens = bm25s.tokenize(query, stemmer=stemmer, stopwords="en")
            
            # Get top results from BM25
            bm25_results, bm25_scores = bm25_retriever.retrieve(
                query_tokens, 
                k=n_results
            )
            
            for i, doc_idx in enumerate(bm25_results[0]):
                if doc_idx < len(bm25_metadata['ids']):
                    doc_id = bm25_metadata['ids'][doc_idx]
                    metadata = bm25_metadata['metadatas'][doc_idx]
                    score = float(bm25_scores[0][i])
                    
                    results_dict[doc_id] = {
                        'id': doc_id,
                        'title': metadata.get('title', 'Untitled'),
                        'excerpt': metadata.get('excerpt', ''),
                        'url': metadata.get('url', ''),
                        'tags': metadata.get('tags', ''),
                        'published_at': metadata.get('published_at', ''),
                        'feature_image': metadata.get('feature_image', ''),
                        'bm25_score': score,
                        'semantic_score': 0.0
                    }
        
        # Semantic search
        if search_type in ["semantic", "hybrid"]:
            semantic_results = semantic_collection.query(
                query_texts=[query],
                n_results=n_results
            )
            
            for i, doc_id in enumerate(semantic_results['ids'][0]):
                metadata = semantic_results['metadatas'][0][i]
                distance = semantic_results['distances'][0][i] if 'distances' in semantic_results else 0
                
                # Convert distance to similarity score
                score = 1.0 / (1.0 + distance)
                
                if doc_id in results_dict:
                    results_dict[doc_id]['semantic_score'] = score
                else:
                    results_dict[doc_id] = {
                        'id': doc_id,
                        'title': metadata.get('title', 'Untitled'),
                        'excerpt': metadata.get('excerpt', ''),
                        'url': metadata.get('url', ''),
                        'tags': metadata.get('tags', ''),
                        'published_at': metadata.get('published_at', ''),
                        'feature_image': metadata.get('feature_image', ''),
                        'bm25_score': 0.0,
                        'semantic_score': score
                    }
        
        # Calculate combined scores
        results = list(results_dict.values())
        
        if search_type == "hybrid":
            # Normalize scores first
            max_bm25 = max([r['bm25_score'] for r in results], default=1.0)
            max_semantic = max([r['semantic_score'] for r in results], default=1.0)
            
            for result in results:
                norm_bm25 = result['bm25_score'] / max_bm25 if max_bm25 > 0 else 0
                norm_semantic = result['semantic_score'] / max_semantic if max_semantic > 0 else 0
                result['score'] = (norm_bm25 * 0.4 + norm_semantic * 0.6)
        elif search_type == "keywords":
            for result in results:
                result['score'] = result['bm25_score']
        else:  # semantic
            for result in results:
                result['score'] = result['semantic_score']
        
        # Sort by score
        results.sort(key=lambda x: x['score'], reverse=True)
        
        # Return top n_results
        return results[:n_results]
        
    except Exception as e:
        print(f"Search error: {e}")
        import traceback
        traceback.print_exc()
        return [{
            'title': 'Error',
            'excerpt': f'Search failed: {str(e)}',
            'url': '',
            'tags': '',
            'published_at': '',
            'feature_image': '',
            'score': 0.0
        }]


def format_results_for_display(results):
    """Format search results for Gradio display."""
    if not results:
        return "No results found."
    
    # Check for error (single result with title "Error")
    if len(results) == 1 and results[0].get('title') == "Error":
        return f"⚠️ {results[0].get('excerpt', 'An error occurred')}"
    
    html_output = '<div style="max-width: 900px; margin: 0 auto;">'
    
    for i, result in enumerate(results, 1):
        # Format published date
        published = result.get('published_at', '')
        if published:
            try:
                from datetime import datetime
                dt = datetime.fromisoformat(published.replace('Z', '+00:00'))
                published = dt.strftime('%B %d, %Y')
            except:
                pass
        
        # Format tags
        tags = result.get('tags', '')
        tags_html = ''
        if tags:
            tag_list = [t.strip() for t in tags.split(',') if t.strip()]
            tags_html = ' '.join([f'<span style="display: inline-block; background: #1237b2; color: white; padding: 2px 8px; border-radius: 3px; font-size: 0.85em; margin-right: 5px;">{tag}</span>' for tag in tag_list])
        
        # Get feature image
        feature_image = result.get('feature_image', '')
        image_html = ''
        if feature_image:
            image_html = f'<img src="{feature_image}" style="width: 100%; height: 200px; object-fit: cover; border-radius: 8px; margin-bottom: 15px;" />'
        
        relevance_score = result.get('score', 0) * 100
        
        html_output += f'''
        <div style="background: white; padding: 25px; margin-bottom: 25px; border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.1); border-left: 4px solid #1237b2;">
            {image_html}
            <div style="display: flex; justify-content: space-between; align-items: start; margin-bottom: 10px;">
                <h3 style="margin: 0; color: #1237b2; font-size: 1.4em;">
                    <a href="{result['url']}" target="_blank" style="color: #1237b2; text-decoration: none;">
                        {i}. {result['title']}
                    </a>
                </h3>
                <span style="background: #f0f4ff; color: #1237b2; padding: 4px 12px; border-radius: 12px; font-size: 0.85em; white-space: nowrap; margin-left: 15px;">
                    {relevance_score:.0f}% match
                </span>
            </div>
            
            {f'<p style="color: #666; font-size: 0.9em; margin: 5px 0 10px 0;">{published}</p>' if published else ''}
            
            {f'<div style="margin-bottom: 10px;">{tags_html}</div>' if tags_html else ''}
            
            <p style="color: #333; line-height: 1.6; margin: 10px 0;">
                {result['excerpt']}
            </p>
            
            <a href="{result['url']}" target="_blank" style="display: inline-block; background: #1237b2; color: white; padding: 8px 16px; text-decoration: none; border-radius: 4px; font-size: 0.9em; margin-top: 10px;">
                Read Full Article →
            </a>
        </div>
        '''
    
    html_output += '</div>'
    return html_output


def create_gradio_interface():
    """Create Gradio interface for search."""
    
    # Custom CSS
    custom_css = """
    .gradio-container {
        max-width: 1200px !important;
        margin: auto !important;
    }
    .header {
        text-align: center;
        padding: 40px 20px;
        background: linear-gradient(135deg, #1237b2 0%, #0e2a7e 100%);
        color: white;
        border-radius: 12px;
        margin-bottom: 30px;
    }
    .header h1 {
        font-size: 2.5em;
        margin-bottom: 10px;
        font-weight: 700;
    }
    .header p {
        font-size: 1.2em;
        opacity: 0.95;
    }
    """
    
    with gr.Blocks(css=custom_css, title="Search") as demo:
        
        # Header
        gr.HTML("""
        <div class="header">
            <h1>🔍 Blueprint Article Search</h1>
            <p>Search through our collection of grassroots project methodologies</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=4):
                query_input = gr.Textbox(
                    label="Search Query",
                    placeholder="e.g., 'community gardens urban planning' or 'cooperative governance models'",
                    lines=1,
                    autofocus=True
                )
            with gr.Column(scale=1):
                search_button = gr.Button("Search", variant="primary", size="lg")
        
        with gr.Accordion("Search Settings", open=False):
            search_type = gr.Radio(
                choices=["keywords", "hybrid", "semantic"],
                value="keywords",
                label="Search Type",
                info="Hybrid combines keyword (BM25) + semantic AI search"
            )
            n_results = gr.Slider(
                minimum=5,
                maximum=100,
                value=10,
                step=1,
                label="Number of Results"
            )
        
        results_output = gr.HTML(label="Search Results")
        
        # Examples
        gr.Examples(
            examples=[
                ["community gardens sustainable food systems", "hybrid", 10],
                ["cooperative ownership models", "semantic", 8],
                ["grassroots organizing strategies", "hybrid", 10],
                ["participatory budgeting governance", "semantic", 8],
                ["urban agriculture resilience", "keywords", 10],
            ],
            inputs=[query_input, search_type, n_results],
        )
        
        # Footer
        gr.HTML("""
        <div style="text-align: center; padding: 30px; color: #666; border-top: 1px solid #e0e0e0; margin-top: 40px;">
            <p>
                <strong>The Blueprint</strong> - Mapping pathways to success for grassroots projects<br>
                <a href="https://the-blueprint.ghost.io" target="_blank" style="color: #1237b2;">Visit our website</a>
            </p>
        </div>
        """)
        
        # Search function wrapper
        def search_and_format(query, search_type, n_results):
            results = search_articles(query, search_type, int(n_results))
            return format_results_for_display(results)
        
        # Event handlers
        search_button.click(
            fn=search_and_format,
            inputs=[query_input, search_type, n_results],
            outputs=results_output
        )
        
        query_input.submit(
            fn=search_and_format,
            inputs=[query_input, search_type, n_results],
            outputs=results_output
        )
    
    return demo


# Initialize index on startup
if __name__ == "__main__":
    print("🚀 Starting Blueprint Search...")
    
    # Download and decrypt index if needed
    if should_update_index() or not CHROMA_DB_PATH.exists():
        print("📥 Initializing search index...")
        download_and_decrypt_index()
    else:
        print("✓ Using cached search index")
    
    # Create and launch interface
    demo = create_gradio_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )