import gradio as gr
import logging
from typing import List, Tuple
import pandas as pd
import os

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Try to import torch and transformers with fallback
try:
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM
    DEPENDENCIES_AVAILABLE = True
except ImportError as e:
    logger.warning(f"Dependencies not available: {e}")
    DEPENDENCIES_AVAILABLE = False
    torch = None
    AutoTokenizer = None
    AutoModelForCausalLM = None

class Qwen3Reranker:
    def __init__(self, model_name="thebajajra/RexReranker-0.6B-v0"):
        if not DEPENDENCIES_AVAILABLE:
            raise ImportError("Required dependencies (torch, transformers) are not available")
        
        self.model_name = model_name
        self.tokenizer = None
        self.model = None
        self.token_false_id = None
        self.token_true_id = None
        self.max_length = 8192
        self.prefix_tokens = None
        self.suffix_tokens = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        self._load_model()
    
    def _load_model(self):
        """Load the tokenizer and model"""
        try:
            logger.info(f"Loading {self.model_name}...")
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name, 
                padding_side='left'
            )
            
            # Load model with appropriate settings
            if torch.cuda.is_available():
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_name,
                    torch_dtype=torch.float16,
                    device_map="auto"
                ).eval()
            else:
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_name
                ).eval()
            
            # Set up tokens
            self.token_false_id = self.tokenizer.convert_tokens_to_ids("no")
            self.token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
            
            # Set up prefix and suffix
            prefix = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n"
            suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
            self.prefix_tokens = self.tokenizer.encode(prefix, add_special_tokens=False)
            self.suffix_tokens = self.tokenizer.encode(suffix, add_special_tokens=False)
            
            logger.info("Model loaded successfully!")
            
        except Exception as e:
            logger.error(f"Error loading model: {e}")
            raise e
    
    def format_instruction(self, instruction: str, query: str, doc: str) -> str:
        """Format the instruction for the reranker"""
        if instruction is None or instruction.strip() == "":
            instruction = 'Given a web search query, retrieve relevant passages that answer the query'
        return f"<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {doc}"
    
    def process_inputs(self, pairs: List[str]) -> dict:
        """Process input pairs for the model"""
        inputs = self.tokenizer(
            pairs,
            padding=True,
            truncation=True,
            return_tensors='pt',
            max_length=self.max_length
        )
        return inputs
    
    def rank_documents(self, instruction: str, query: str, documents: List[str]) -> List[Tuple[str, float, str]]:
        """Rank documents based on their relevance to the query"""
        if not DEPENDENCIES_AVAILABLE:
            return [(doc[:100] + "...", 0.5, "Dependencies not available") for doc in documents]
        
        results = []
        
        for i, doc in enumerate(documents):
            try:
                # Format the instruction
                formatted_instruction = self.format_instruction(instruction, query, doc)
                
                # Tokenize
                inputs = self.tokenizer(
                    formatted_instruction,
                    return_tensors='pt',
                    max_length=self.max_length,
                    truncation=True
                )
                
                if torch.cuda.is_available():
                    inputs = {k: v.cuda() for k, v in inputs.items()}
                
                # Get model output
                with torch.no_grad():
                    outputs = self.model(**inputs)
                    logits = outputs.logits[0, -1, :]
                    
                    # Get probabilities for "yes" and "no" tokens
                    yes_prob = torch.softmax(logits, dim=-1)[self.token_true_id].item()
                    no_prob = torch.softmax(logits, dim=-1)[self.token_false_id].item()
                    
                    # Calculate relevance score (probability of "yes")
                    relevance_score = yes_prob / (yes_prob + no_prob)
                
                # Truncate document for display
                display_doc = doc[:200] + "..." if len(doc) > 200 else doc
                results.append((display_doc, relevance_score, f"Document {i+1}"))
                
            except Exception as e:
                logger.error(f"Error processing document {i+1}: {e}")
                display_doc = doc[:200] + "..." if len(doc) > 200 else doc
                results.append((display_doc, 0.0, f"Error: {str(e)[:50]}..."))
        
        # Sort by relevance score (highest first)
        results.sort(key=lambda x: x[1], reverse=True)
        return results

# Initialize the reranker
try:
    reranker = Qwen3Reranker()
except Exception as e:
    logger.error(f"Failed to initialize reranker: {e}")
    reranker = None

def rerank_documents(instruction, query, documents_text):
    """Gradio interface function"""
    if not reranker:
        return pd.DataFrame([["Error", "Model not loaded", 0.0]], 
                          columns=["Document", "Relevance Score", "Rank"])
    
    if not query.strip():
        return pd.DataFrame([["Error", "Please provide a query", 0.0]], 
                          columns=["Document", "Relevance Score", "Rank"])
    
    if not documents_text.strip():
        return pd.DataFrame([["Error", "Please provide documents", 0.0]], 
                          columns=["Document", "Relevance Score", "Rank"])
    
    # Split documents by double newlines or numbered list format
    documents = []
    if '\n\n' in documents_text:
        documents = [doc.strip() for doc in documents_text.split('\n\n') if doc.strip()]
    else:
        # Try to split by numbered format (1., 2., etc.)
        lines = documents_text.strip().split('\n')
        current_doc = ""
        for line in lines:
            if line.strip() and (line.strip()[0].isdigit() and '.' in line[:5]):
                if current_doc:
                    documents.append(current_doc.strip())
                current_doc = line
            else:
                current_doc += "\n" + line
        if current_doc:
            documents.append(current_doc.strip())
    
    if not documents:
        documents = [documents_text]  # Treat as single document
    
    # Rank documents
    results = reranker.rank_documents(instruction, query, documents)
    
    # Create DataFrame for display
    df_data = []
    for i, (doc, score, label) in enumerate(results):
        df_data.append([f"#{i+1}", doc, f"{score:.4f}"])
    
    return pd.DataFrame(df_data, columns=["Rank", "Document", "Relevance Score"])

def create_gradio_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(title="RexRanker: E-commerce reranker", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🔍 RexRanker: E-commerce reranker
        
        This tool uses the **RexReranker-0.6B** model to rank products by their relevance to search query.
        
        ## How to use:
        1. **Instruction** (optional): Provide context for the ranking task
        2. **Query**: Enter your search query
        3. **Documents**: Enter multiple documents separated by double newlines (\\n\\n) or as a numbered list
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                instruction_input = gr.Textbox(
                    label="Instruction (Optional)",
                    placeholder="Given a web search query, retrieve relevant passages that answer the query",
                    value="Given a web search query, retrieve relevant passages that answer the query",
                    lines=2
                )
                
                query_input = gr.Textbox(
                    label="Query",
                    placeholder="Enter your search query here...",
                    lines=2
                )
                
                documents_input = gr.Textbox(
                    label="Documents to Rank",
                    placeholder="Enter documents separated by double newlines...\n\nDocument 1 content here\n\nDocument 2 content here\n\nDocument 3 content here",
                    lines=10
                )
                
                rank_button = gr.Button("🔍 Rank Documents", variant="primary")
                
                gr.Markdown("### Example:")
                gr.Examples(
                    examples=[
                        [
                            "Given a e-commerce search query, retrieve relevant products that are relevant to the query",
                            "set of 10 blue off road bulbs",
                            "Set of 10 Blue Light Bulbs Ten 3157 Pack Replacement Off Road AutoAndArt"
                        ]
                    ],
                    inputs=[instruction_input, query_input, documents_input]
                )
            
            with gr.Column(scale=1):
                # FIXED: Remove the height parameter
                results_display = gr.DataFrame(
                    label="Ranking Results",
                    headers=["Rank", "Document", "Relevance Score"],
                    interactive=False
                )
        
        rank_button.click(
            fn=rerank_documents,
            inputs=[instruction_input, query_input, documents_input],
            outputs=[results_display]
        )
        
        gr.Markdown("""
        ### About the Model
        - **Model**: thebajajra/RexReranker-0.6B
        - **Task**: Document reranking based on query relevance
        - **Output**: Relevance scores between 0 and 1 (higher = more relevant)
        """)
    
    return demo

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