import os
os.environ['ANONYMIZED_TELEMETRY'] = 'False'

import zipfile
import chromadb
from sentence_transformers import SentenceTransformer
import gradio as gr
from fastapi import FastAPI
from pydantic import BaseModel
import re

# Extract and load database
DB_PATH = "./medqa_db"
if not os.path.exists(DB_PATH) and os.path.exists("./medqa_db.zip"):
    print("📦 Extracting database...")
    with zipfile.ZipFile("./medqa_db.zip", 'r') as z:
        z.extractall(".")
    print("✅ Database extracted")

print("🔌 Loading ChromaDB...")
client = chromadb.PersistentClient(path=DB_PATH)
collection = client.get_collection("medqa")
print(f"✅ Loaded {collection.count()} questions")

print("🧠 Loading MedCPT model...")
model = SentenceTransformer('ncbi/MedCPT-Query-Encoder')
print("✅ Model ready")

# ============================================================================
# Deduplication function
# ============================================================================
def deduplicate_results(results, target_count):
    """
    Remove duplicate questions based on:
    1. High text similarity (>0.92) - catches near-exact duplicates
    2. Same answer + moderate similarity (>0.85) - catches conceptual duplicates
    """
    if not results['documents'][0]:
        return results
    
    documents = results['documents'][0]
    metadatas = results['metadatas'][0]
    distances = results['distances'][0]
    
    selected_indices = []
    
    for i in range(len(documents)):
        is_duplicate = False
        current_answer = metadatas[i].get('answer', '')
        
        for j in selected_indices:
            selected_answer = metadatas[j].get('answer', '')
            dist_diff = abs(distances[i] - distances[j])
            
            if dist_diff < 0.08:
                is_duplicate = True
                break
            
            if current_answer == selected_answer and dist_diff < 0.15:
                is_duplicate = True
                break
        
        if not is_duplicate:
            selected_indices.append(i)
            
            if len(selected_indices) >= target_count:
                break
    
    return {
        'documents': [[documents[i] for i in selected_indices]],
        'metadatas': [[metadatas[i] for i in selected_indices]],
        'distances': [[distances[i] for i in selected_indices]],
        'ids': [[results['ids'][0][i] for i in selected_indices]] if 'ids' in results else None
    }

# ============================================================================
# Search function with deduplication
# ============================================================================
def search(query, num_results=3, source_filter=None):
    emb = model.encode(query).tolist()
    
    where_clause = None
    if source_filter and source_filter != "all":
        where_clause = {"source": source_filter}
    
    fetch_count = min(num_results * 4, 50)
    
    results = collection.query(
        query_embeddings=[emb], 
        n_results=fetch_count,
        where=where_clause
    )
    
    return deduplicate_results(results, num_results)

# ============================================================================
# Parser to extract question structure
# ============================================================================
def parse_question_document(doc_text, metadata):
    """Extract question and choices from document text - NO TRUNCATION."""
    
    lines = doc_text.split('\n')
    question_lines = []
    options_started = False
    options = {}
    
    for line in lines:
        line = line.strip()
        if not line:
            continue
            
        # Check if this is an option line (A., B., C., etc.)
        option_match = re.match(r'^([A-E])[\.\)]\s*(.+)$', line)
        
        if option_match:
            options_started = True
            letter = option_match.group(1)
            text = option_match.group(2).strip()
            options[letter] = text
        elif not options_started:
            question_lines.append(line)
    
    # Reconstruct FULL question text - no truncation
    question_text = ' '.join(question_lines).strip()
    
    answer_idx = metadata.get('answer_idx', 'N/A')
    answer_text = metadata.get('answer', 'N/A')
    
    # If answer_text is just the letter, map it to the actual option text
    if answer_text in options:
        answer_text = options[answer_text]
    
    return {
        'question': question_text,
        'choices': options,
        'correct_answer_letter': answer_idx,
        'correct_answer_text': answer_text
    }

# ============================================================================
# Enhanced Gradio UI
# ============================================================================
def ui_search(query, num_results=3, source_filter="all"):
    if not query.strip():
        return "💡 Enter a medical query to search"
    
    try:
        r = search(query, num_results, source_filter if source_filter != "all" else None)
        
        if not r['documents'][0]:
            return "❌ No results found"
        
        out = f"🔍 Found {len(r['documents'][0])} unique results\n\n"
        
        for i in range(len(r['documents'][0])):
            source = r['metadatas'][0][i].get('source', 'unknown')
            distance = r['distances'][0][i]
            similarity = 1 - distance
            
            # Source emoji
            if source == 'medgemini':
                source_icon = "🔬"
                source_name = "Med-Gemini"
            elif source.startswith('medqa_'):
                source_icon = "📚"
                split = source.replace('medqa_', '').upper()
                source_name = f"MedQA {split}"
            else:
                source_icon = "📄"
                source_name = source.upper()
            
            out += f"\n{'='*70}\n"
            out += f"{source_icon} Result {i+1} | {source_name} | Similarity: {similarity:.3f}\n"
            out += f"{'='*70}\n\n"
            out += r['documents'][0][i]
            
            answer = r['metadatas'][0][i].get('answer', 'N/A')
            out += f"\n\n✅ CORRECT ANSWER: {answer}\n"
            
            explanation = r['metadatas'][0][i].get('explanation', '')
            if explanation and explanation.strip():
                out += f"\n💡 EXPLANATION:\n{explanation}\n"
            
            out += "\n"
        
        return out
    
    except Exception as e:
        return f"❌ Error: {e}"

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), title="MedQA Search") as demo:
    gr.Markdown("""
    # 🏥 MedQA Semantic Search
    
    Search across **Med-Gemini** (expert explanations) and **MedQA** (USMLE questions) databases.
    Uses medical-specific embeddings (MedCPT) for accurate retrieval.
    
    ✨ **Features**: Automatic deduplication, structured output for AI integration
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            query_input = gr.Textbox(
                label="Medical Query",
                placeholder="e.g., hyponatremia, myocardial infarction, diabetes management...",
                lines=2
            )
        with gr.Column(scale=1):
            num_results = gr.Slider(
                minimum=1,
                maximum=10,
                value=3,
                step=1,
                label="Number of Results"
            )
    
    with gr.Row():
        source_filter = gr.Radio(
            choices=["all", "medgemini", "medqa_train", "medqa_dev", "medqa_test"],
            value="all",
            label="Filter by Source"
        )
    
    search_btn = gr.Button("🔍 Search", variant="primary", size="lg")
    
    output = gr.Textbox(
        label="Search Results",
        lines=25,
        max_lines=50
    )
    
    search_btn.click(
        fn=ui_search,
        inputs=[query_input, num_results, source_filter],
        outputs=output
    )
    
    query_input.submit(
        fn=ui_search,
        inputs=[query_input, num_results, source_filter],
        outputs=output
    )
    
    gr.Markdown("""
    ### 📊 Database Info
    
    **Med-Gemini**: Expert-relabeled questions with detailed explanations  
    **MedQA**: USMLE-style questions (Train/Dev/Test splits)
    
    **Total Questions**: ~10,000+ USMLE-style questions
    """)
    
    gr.Examples(
        examples=[
            ["hyponatremia", 3, "all"],
            ["myocardial infarction treatment", 2, "medgemini"],
            ["diabetes complications", 3, "all"],
            ["antibiotics for pneumonia", 2, "medqa_train"]
        ],
        inputs=[query_input, num_results, source_filter]
    )

# ============================================================================
# FastAPI with structured JSON output (for OpenAI integration)
# ============================================================================
app = FastAPI()

class SearchRequest(BaseModel):
    query: str
    num_results: int = 3
    source_filter: str = None

@app.post("/search_medqa")
def api_search(req: SearchRequest):
    """
    Search MedQA and return structured exemplars.
    Returns COMPLETE question text with no truncation.
    """
    
    r = search(req.query, req.num_results, req.source_filter)
    
    if not r['documents'][0]:
        return {"results": []}
    
    results = []
    for i in range(len(r['documents'][0])):
        doc_text = r['documents'][0][i]
        metadata = r['metadatas'][0][i]
        
        # Parse the document into structured format
        parsed = parse_question_document(doc_text, metadata)
        
        # Build complete result object
        result = {
            "result_number": i + 1,
            "question": parsed['question'],  # FULL question text
            "choices": parsed['choices'],
            "correct_answer": parsed['correct_answer_letter'],
            "correct_answer_text": parsed['correct_answer_text'],
            "explanation": metadata.get('explanation', ''),
            "has_explanation": bool(metadata.get('explanation', '').strip()),
            "source": metadata.get('source', 'unknown'),
            "exam_type": metadata.get('exam_type', 'unknown'),
            "split": metadata.get('split', 'unknown'),
            "similarity": round(1 - r['distances'][0][i], 3),
            "metamap_phrases": metadata.get('metamap_phrases', '')
        }
        
        results.append(result)
    
    return {"results": results}

app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)