import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModel
import numpy as np
from typing import List, Dict
import pandas as pd
import os
from pinecone import Pinecone


# =========================
# Retriever Class
# =========================
class ParrotletRetriever:
    def __init__(self, model_name: str, pinecone_api_key: str, index_name: str):
        """Initialize model and Pinecone client."""
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"🚀 Loading model on {self.device}...")

        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
        self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True, token=HF_TOKEN)
        self.model.to(self.device)
        self.model.eval()

        # Pinecone initialization (unchanged)
        self.pinecone_namespace = os.environ.get("NAMESPACE")
        self.pinecone_client = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))
        self.pinecone_index = self.pinecone_client.Index(host=os.environ.get("PINECONE_HOST"))

        print(f"Connected to Pinecone index: {index_name}")
        if self.pinecone_namespace:
            print(f"🔹 Using namespace: {self.pinecone_namespace}")

    # --------------------------
    # Mean Pooling
    # --------------------------
    def mean_pooling(self, model_output, attention_mask):
        """Mean pooling for sentence embeddings."""
        token_embeddings = model_output[0]
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
            input_mask_expanded.sum(1), min=1e-9
        )

    # --------------------------
    # Text Encoder
    # --------------------------
    def encode(self, texts: List[str]) -> np.ndarray:
        """Encode text into normalized embeddings."""
        with torch.no_grad():
            encoded_input = self.tokenizer(
                texts, padding=True, truncation=True, max_length=512, return_tensors="pt"
            ).to(self.device)

            model_output = self.model(**encoded_input)
            embeddings = self.mean_pooling(model_output, encoded_input["attention_mask"])
            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
            return embeddings.cpu().numpy()

    # --------------------------
    # Pinecone Search
    # --------------------------
    def search(self, query: str, top_k: int = 5) -> List[Dict]:
        """Search Pinecone index."""
        query_vector = self.encode([query])[0]

        results = self.pinecone_index.query(
            namespace=self.pinecone_namespace,
            vector=query_vector.tolist(),
            top_k=top_k,
            include_metadata=True,
            include_values=False,
        )

        docs = []
        for i, match in enumerate(results["matches"]):
            text = match["metadata"].get("text", "[No text metadata]")
            docs.append({
                "Rank": i + 1,
                "Score": f"{match['score']:.4f}",
                "Document": text
                "ID": match["concept_id"],
            })

        return docs



MODEL_NAME = "ekacare/parrotlet-e"
retriever = ParrotletRetriever(MODEL_NAME, PINECONE_API_KEY, PINECONE_INDEX_NAME)


def retrieve_documents(query: str, top_k: int = 5):
    """Perform retrieval and return results."""
    if not query.strip():
        return pd.DataFrame(), "⚠️ Please enter a valid query."

    try:
        results = retriever.search(query, top_k)
        if not results:
            return pd.DataFrame(), "⚠️ No results found."

        df = pd.DataFrame(results)
        status = f"✅ Retrieved top {len(results)} documents."
        return df, status

    except Exception as e:
        return pd.DataFrame(), f"❌ Error: {str(e)}"


# =========================
# Gradio Interface
# =========================
SAMPLE_QUERIES = [
    "dm t2",
    "छाती में दर्द",
    "talenovu",
    "వాంతులు"
]

with gr.Blocks(title="Parrotlet-E Retrieval (Pinecone)", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🦜 Parrotlet-E: Indic Medical Entities Retrieval
    Retrieve top-5 semantically similar medical entities from Pinecone using Parrotlet-E embeddings.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            query_input = gr.Textbox(
                label="Enter a medical term",
                placeholder="Type your query here...",
                lines=1,
            )

            examples = gr.Examples(
                examples=SAMPLE_QUERIES,
                inputs=query_input,
                label="Example Queries"
            )

            top_k_slider = gr.Slider(
                minimum=1,
                maximum=10,
                value=5,
                step=1,
                label="Top-K Retrieved"
            )

            search_btn = gr.Button("🔍 Retrieve", variant="primary")

        with gr.Column(scale=2):
            results_output = gr.Dataframe(
                headers=["Rank", "Score", "Document", "ID"],
                datatype=["number", "str", "str", "str"],
                interactive=False,
                wrap=True
            )
            status_box = gr.Textbox(label="Status", interactive=False)

    search_btn.click(
        fn=retrieve_documents,
        inputs=[query_input, top_k_slider],
        outputs=[results_output, status_box],
    )


# =========================
# Run App
# =========================
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)