import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from utils import get_embed_model


# ============================================================================
# LAZY-LOAD RERANKER AND QUERY REWRITER (load on first use)
# ============================================================================

_rerank_tokenizer = None
_rerank_model = None
_rewritter_llm = None

def get_reranker():
    """Load reranker model (lazy-loaded on first use)"""
    global _rerank_tokenizer, _rerank_model
    if _rerank_tokenizer is None:
        _rerank_tokenizer = AutoTokenizer.from_pretrained("castorini/monot5-base-msmarco")
        _rerank_model = AutoModelForSeq2SeqLM.from_pretrained("castorini/monot5-base-msmarco")
        _rerank_model.eval()
    return _rerank_tokenizer, _rerank_model

def get_query_rewriter():
    """Load query rewriter (lazy-loaded on first use)"""
    global _rewritter_llm
    if _rewritter_llm is None:
        _rewritter_llm = pipeline(
            "text2text-generation",
            model="google/flan-t5-small",
            max_length=64,
            do_sample=False,
            temperature=0.3,
            repetition_penalty=1.3,
            no_repeat_ngram_size=2
        )
    return _rewritter_llm


def rewrite_query(user_query):
    """Rewrite user query to be more specific and medical-focused.
    Returns: str(rewritten_query)"""
    
    rewritter_llm = get_query_rewriter()
    
    prompt = f"""Rewrite the input into a clear medical question following these patterns
    Input: my head hurts
    Output: What causes headaches?

    Input: i keep vomiting but feel ok afterwards
    Output: What causes cyclic vomiting?

    Input: chest pain when breathing
    Output: What causes chest pain during breathing?

    Input: {user_query}
    Output:
    """
    
    llm_output = rewritter_llm(prompt)
    rewritten_query = llm_output[0]['generated_text']
    rewritten_query = rewritten_query.replace("Output:", "").strip()
    
    return rewritten_query.strip()


def get_monot5_scores(rewritten_query, chunks):
    """Calculate relevance scores for chunks against query using MonoT5 reranker.
    Returns: list(scores)"""
    
    rerank_tokenizer, rerank_model = get_reranker()
    
    # Format input as "Query: ABC? Document: XYZ... Relevant:"
    inputs = [f"Query: {rewritten_query} Document: {c} Relevant:" for c in chunks]
    
    # Tokenize
    tokenized_inputs = rerank_tokenizer(
        inputs,
        padding=True,
        truncation=True,
        return_tensors='pt'
    )
    
    # Generate predictions
    with torch.no_grad():
        outputs = rerank_model.generate(
            input_ids=tokenized_inputs['input_ids'],
            attention_mask=tokenized_inputs['attention_mask'],
            max_new_tokens=1,
            return_dict_in_generate=True,
            output_scores=True
        )
        
        # Extract "true" token probability scores
        true_token_id = rerank_tokenizer.encode("true")[0]
        batch_scores = outputs.scores[0][:, true_token_id]
        
        return batch_scores.tolist()


def retrieve(vector_store, query, top_k=3, use_reranking=True, detail=False):
    """Retrieve relevant context chunks for a query with optional reranking.
    Returns: list(contexts)"""
    
    embed_model = get_embed_model()
    
    is_single_query = isinstance(query, str)
    queries = [query] if is_single_query else query
    
    # Rewrite queries for better matching
    rewritten_queries = [rewrite_query(q) for q in queries]
    
    # Embed rewritten queries
    q_embeddings = embed_model.encode(rewritten_queries).tolist()
    
    # Vector search
    search_results = vector_store.query(
        query_embeddings=q_embeddings,
        n_results=10
    )
    
    all_contexts = []
    
    for i in range(len(queries)):
        contexts_for_query = []
        
        if not search_results['documents'][i]:
            all_contexts.append([])
            continue
        
        retrieved_chunks = search_results['documents'][i]
        retrieved_metas = search_results['metadatas'][i]
        retrieved_ids = search_results['ids'][i]
        retrieved_distances = search_results.get('distances', [[]])[i]
        
        # Rerank if enabled
        if use_reranking:
            rerank_scores = get_monot5_scores(rewritten_queries[i], retrieved_chunks)
            
            candidates = []
            for r in range(len(retrieved_chunks)):
                candidates.append({
                    'chunk_id': retrieved_ids[r],
                    'chunk_answer': retrieved_chunks[r],
                    'question': retrieved_metas[r]['question'],
                    'focus_area': retrieved_metas[r].get('focus_area', 'Unknown'),
                    'chroma_dist': retrieved_distances[r],
                    'rerank_score': rerank_scores[r]
                })
            
            # Sort by rerank score
            candidates.sort(key=lambda s: s['rerank_score'], reverse=True)
            sorted_results = candidates[:top_k]
        else:
            # Just use top-k from vector search
            sorted_results = [
                {
                    'chunk_id': retrieved_ids[r],
                    'chunk_answer': retrieved_chunks[r],
                    'question': retrieved_metas[r]['question'],
                    'focus_area': retrieved_metas[r].get('focus_area', 'Unknown'),
                    'chroma_dist': retrieved_distances[r]
                }
                for r in range(min(top_k, len(retrieved_chunks)))
            ]
        
        # Format output
        for item_dict in sorted_results:
            item = {
                'chunk_id': item_dict['chunk_id'],
                'chunk_answer': item_dict['chunk_answer'],
                'question': item_dict['question']
            }
            if detail:
                item['focus_area'] = item_dict['focus_area']
                item['chroma_distance'] = round(item_dict['chroma_dist'], 3)
                if 'rerank_score' in item_dict:
                    item['rerank_score'] = round(item_dict['rerank_score'], 3)
            contexts_for_query.append(item)
        
        all_contexts.append(contexts_for_query)
    
    if is_single_query:
        return all_contexts[0]
    else:
        return all_contexts