import os
import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer
import pickle
import gradio as gr
import gdown
import requests
from tqdm import tqdm
import time
import re
import traceback
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, 
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    handlers=[logging.StreamHandler()])
logger = logging.getLogger("SportsChatbot")

def initialize_llm():
    from transformers import AutoModelForCausalLM, AutoTokenizer
    
    model_id = "deepseek-ai/deepseek-coder-1.3b-instruct"
    model_dir = "models/deepseek-coder-1.3b-instruct"
    
    if not os.path.exists(model_dir):
        os.makedirs(model_dir, exist_ok=True)
        logger.info(f"Downloading DeepSeek Coder model to {model_dir}...")
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float32,  # Use float32 for CPU
            low_cpu_mem_usage=True,
        )
        model.save_pretrained(model_dir)
        tokenizer.save_pretrained(model_dir)
    
    # Fix for attention mask and pad token warnings
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    # Explicitly set pad_token_id if it's not set
    if tokenizer.pad_token_id is None:
        logger.info("Setting pad_token_id to eos_token_id")
        tokenizer.pad_token_id = tokenizer.eos_token_id
    
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True
    )
    
    logger.info(f"Model loaded successfully from {model_dir}")
    return model, tokenizer

def download_file_with_progress(url: str, filename: str):
    """Download a file with progress bar using requests"""
    response = requests.get(url, stream=True)
    total_size = int(response.headers.get('content-length', 0))
    
    with open(filename, 'wb') as file, tqdm(
        desc=filename,
        total=total_size,
        unit='iB',
        unit_scale=True,
        unit_divisor=1024,
    ) as progress_bar:
        for data in response.iter_content(chunk_size=1024):
            size = file.write(data)
            progress_bar.update(size)

class SentenceTransformerRetriever:
    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", cache_dir: str = "embeddings_cache"):
        self.device = torch.device("cpu")
        self.model = SentenceTransformer(model_name, device=str(self.device))
        self.doc_embeddings = None
        self.cache_dir = cache_dir
        os.makedirs(cache_dir, exist_ok=True)
        
    def load_specific_cache(self, cache_filename: str, drive_link: str) -> dict:
        cache_path = os.path.join(self.cache_dir, cache_filename)
        
        if not os.path.exists(cache_path):
            logger.info(f"Cache file not found. Downloading from Google Drive...")
            try:
                gdown.download(drive_link, cache_path, quiet=False)
            except Exception as e:
                logger.error(f"Failed to download cache file: {str(e)}")
                raise Exception(f"Failed to download cache file: {str(e)}")
            
            if not os.path.exists(cache_path):
                raise FileNotFoundError(f"Failed to download cache file to {cache_path}")
        
        logger.info(f"Loading cache from: {cache_path}")
        with open(cache_path, 'rb') as f:
            return pickle.load(f)

    def encode(self, texts: list) -> torch.Tensor:
        embeddings = self.model.encode(texts, convert_to_tensor=True, show_progress_bar=True)
        return F.normalize(embeddings, p=2, dim=1)

    def store_embeddings(self, embeddings: torch.Tensor):
        self.doc_embeddings = embeddings

    def search(self, query_embedding: torch.Tensor, k: int):
        if self.doc_embeddings is None:
            raise ValueError("No document embeddings stored!")

        similarities = F.cosine_similarity(query_embedding, self.doc_embeddings)
        scores, indices = torch.topk(similarities, k=min(k, similarities.shape[0]))
        
        return indices.cpu(), scores.cpu()

class TextProcessor:
    @staticmethod
    def extract_relevant_info(documents, query):
        """Extract only the most relevant parts from documents based on the query"""
        # Simple keyword-based extraction
        query_tokens = set(re.findall(r'\b\w+\b', query.lower()))
        relevant_parts = []
        
        for doc in documents:
            # Split document into sentences
            sentences = re.split(r'(?<=[.!?])\s+', doc)
            
            for sentence in sentences:
                # Count matching keywords
                sentence_tokens = set(re.findall(r'\b\w+\b', sentence.lower()))
                overlap = len(query_tokens.intersection(sentence_tokens))
                
                # Keep sentences with sufficient keyword overlap or short enough to be important
                if overlap > 0 or len(sentence) < 100:
                    relevant_parts.append(sentence)
        
        # Join the relevant parts
        result = " ".join(relevant_parts)
        
        # If result is too long, truncate more aggressively
        if len(result) > 1000:  # Reduced from 2000 to 1000
            result = result[-1000:]
            
        return result

def truncate_to_token_limit(text, tokenizer, max_tokens=500):
    """Truncate text to fit within token limit"""
    tokens = tokenizer.encode(text)
    if len(tokens) <= max_tokens:
        return text
    
    # Truncate tokens and decode back to text
    truncated_tokens = tokens[-max_tokens:]
    return tokenizer.decode(truncated_tokens)

def generate_response(model, tokenizer, prompt):
    """Generate a response using the DeepSeek model"""
    logger.info("Tokenizing input prompt")
    inputs = tokenizer(prompt, return_tensors="pt", padding=True)  # Add padding=True to ensure attention_mask is set
    input_length = inputs["input_ids"].shape[1]
    
    # Print prompt length for debugging
    logger.info(f"Prompt length: {input_length} tokens")
    
    logger.info("Starting model generation")
    try:
        with torch.no_grad():
            output = model.generate(
                inputs["input_ids"],
                attention_mask=inputs["attention_mask"],  # Pass attention_mask
                max_new_tokens=256,  # Generate up to 256 new tokens
                temperature=0.1,
                top_p=0.1,
                top_k=10,
                repetition_penalty=1.1,
                do_sample=True,
                num_return_sequences=1,
                pad_token_id=tokenizer.pad_token_id,  # Set pad_token_id explicitly
            )
        
        logger.info("Model generation completed successfully")
        response = tokenizer.decode(output[0], skip_special_tokens=True)
        # Extract just the response part (after the prompt)
        generation = response[len(tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True)):].strip()
        logger.info(f"Generated response length: {len(generation)} characters")
        return generation
    except Exception as e:
        logger.error(f"Error during generation: {str(e)}")
        logger.error(traceback.format_exc())
        return f"Error during text generation: {str(e)}"

def process_query(query: str, cache_data: dict, model, tokenizer) -> str:
    try:
        logger.info(f"Processing query: {query}")
        retriever = SentenceTransformerRetriever()
        retriever.store_embeddings(cache_data['embeddings'])
        documents = cache_data['documents']
        
        # Get relevant document indices
        logger.info("Encoding query and searching for relevant documents")
        query_embedding = retriever.encode([query])
        indices, scores = retriever.search(query_embedding, k=5)
        
        # Get only documents with score above threshold
        selected_indices = [idx for idx, score in zip(indices.tolist(), scores.tolist()) if score > 0.3]
        logger.info(f"Found {len(selected_indices)} relevant documents with score > 0.3")
        
        if not selected_indices:
            return "I don't have enough information to answer that question accurately."
        
        relevant_docs = [documents[idx] for idx in selected_indices]
        
        # Extract only the most relevant information from documents
        logger.info("Extracting relevant information from documents")
        extracted_text = TextProcessor.extract_relevant_info(relevant_docs, query)
        
        # Ensure we don't exceed token limits
        extracted_text = truncate_to_token_limit(extracted_text, tokenizer, max_tokens=500)
        
        if not extracted_text:
            return "I couldn't find relevant information to answer your question."
            
        # Format prompt for DeepSeek model
        prompt = f"""<问题>
Use the following information to answer the sports question. Only use facts mentioned in the given information.

Information:
{extracted_text}

Question: {query}
</问题>

<回答>"""

        logger.info("Generating response")
        start_time = time.time()
        response = generate_response(model, tokenizer, prompt)
        elapsed_time = time.time() - start_time
        logger.info(f'Inference Time: {elapsed_time:.2f} seconds')
        
        return response
    
    except Exception as e:
        logger.error(f"Error processing query: {str(e)}")
        logger.error(traceback.format_exc())
        return f"I encountered an error while processing your question: {str(e)}"

class SportsChatbot:
    def __init__(self):
        self.cache_filename = "embeddings_2296.pkl"
        self.cache_drive_link = "https://drive.google.com/uc?id=1LuJdnwe99C0EgvJpyfHYCKzUvj94FWlC"
        self.cache_data = None
        self.model = None
        self.tokenizer = None
        self.initialize_pipeline()

    def initialize_pipeline(self):
        try:
            logger.info("Initializing SportsChatbot pipeline")
            # Initialize retriever and load cache
            retriever = SentenceTransformerRetriever()
            self.cache_data = retriever.load_specific_cache(self.cache_filename, self.cache_drive_link)
            
            # Initialize DeepSeek model
            self.model, self.tokenizer = initialize_llm()
            logger.info("SportsChatbot pipeline initialized successfully")
            
        except Exception as e:
            logger.error(f"Error initializing the application: {str(e)}")
            logger.error(traceback.format_exc())
            raise Exception(f"Error initializing the application: {str(e)}")

    def process_question(self, question: str, progress=gr.Progress()):
        try:
            if not question.strip():
                return "Please enter a question!"
            
            logger.info(f"Received question: {question}")
            progress(0.1, desc="Processing query...")
            
            # Add intermediate progress steps for better feedback
            progress(0.3, desc="Searching knowledge base...")
            progress(0.5, desc="Analyzing relevant information...")
            response = process_query(question, self.cache_data, self.model, self.tokenizer)
            progress(0.9, desc="Finalizing response...")
            
            logger.info("Response generated successfully")
            return response
        except Exception as e:
            logger.error(f"Error in process_question: {str(e)}")
            logger.error(traceback.format_exc())
            return f"Sorry, an error occurred: {str(e)}"

def create_demo():
    try:
        logger.info("Creating Gradio demo")
        chatbot = SportsChatbot()
        
        with gr.Blocks(title="The Sport Chatbot") as demo:
            gr.Markdown("# The Sport Chatbot")
            gr.Markdown("### Using ESPN API")
            gr.Markdown("Hey there! 👋 I can help you with information on Ice Hockey, Baseball, American Football, Soccer, and Basketball. With access to the ESPN API, I'm up to date with the latest details for these sports up until October 2024.")
            gr.Markdown("Got any general questions? Feel free to ask—I'll do my best to provide answers based on the information I've been trained on!")
            
            with gr.Row():
                question_input = gr.Textbox(
                    label="Enter your question:",
                    placeholder="Type your sports-related question here...",
                    lines=2
                )
            
            with gr.Row():
                submit_btn = gr.Button("Get Answer", variant="primary")
            
            # Add a debug output area
            with gr.Row():
                debug_output = gr.Markdown(label="Debug Info", visible=True)
            
            with gr.Row():
                answer_output = gr.Markdown(label="Answer")
                
            # Define a debug function to test button clicks
            def update_debug_info(question):
                logger.info(f"Button clicked with question: {question}")
                return f"Processing question: '{question}' at {time.strftime('%H:%M:%S')}"
            
            # Connect the debug function to show activity
            submit_btn.click(
                fn=update_debug_info,
                inputs=question_input,
                outputs=debug_output,
            ).then(
                fn=chatbot.process_question,
                inputs=question_input,
                outputs=answer_output,
                api_name="answer_question"
            )
            
            gr.Examples(
                examples=[
                    "Who won the NBA championship in 2023?",
                    "What are the basic rules of ice hockey?",
                    "Tell me about the NFL playoffs format.",
                ],
                inputs=question_input,
            )

        logger.info("Gradio demo created successfully")
        return demo
    except Exception as e:
        logger.error(f"Error creating demo: {str(e)}")
        logger.error(traceback.format_exc())
        raise

if __name__ == "__main__":
    try:
        logger.info("Starting Sports Chatbot application")
        demo = create_demo()
        logger.info("Launching Gradio interface")
        demo.launch()
    except Exception as e:
        logger.error(f"Application startup failed: {str(e)}")
        logger.error(traceback.format_exc())