import os
import warnings
warnings.filterwarnings("ignore", category=UserWarning)

import streamlit as st
import torch
import torch.nn.functional as F
import re
import requests
from embedding_processor import SentenceTransformerRetriever, process_data
import pickle
import logging
import sys
from llama_cpp import Llama
from tqdm import tqdm

# Set page config first
st.set_page_config(
    page_title="The Sport Chatbot",
    page_icon="🏆",
    layout="wide"
)

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

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)

@st.cache_data
def load_from_drive(file_id: str):
    """Load pickle file directly from Google Drive"""
    try:
        url = f"https://drive.google.com/uc?id={file_id}&export=download"
        session = requests.Session()
        response = session.get(url, stream=True)
        
        for key, value in response.cookies.items():
            if key.startswith('download_warning'):
                url = f"{url}&confirm={value}"
                response = session.get(url, stream=True)
                break
        
        content = response.content
        print(f"Successfully downloaded {len(content)} bytes")
        return pickle.loads(content)
        
    except Exception as e:
        print(f"Detailed error: {str(e)}")
        st.error(f"Error loading file from Drive: {str(e)}")
        return None

# @st.cache_resource(show_spinner=False)
# def load_llama_model():
#     """Load Llama model with caching"""
#     try:
#         model_path = "mistral-7b-v0.1.Q4_K_M.gguf"
        
#         if not os.path.exists(model_path):
#             st.info("Downloading model... This may take a while.")
#             direct_url = "https://huggingface.co/TheBloke/Mistral-7B-v0.1-GGUF/resolve/main/mistral-7b-v0.1.Q4_K_M.gguf"
#             download_file_with_progress(direct_url, model_path)

#         llm_config = {
#             "model_path": model_path,
#             "n_ctx": 2048,
#             "n_threads": 4,
#             "n_batch": 512,
#             "n_gpu_layers": 0,
#             "verbose": False
#         }
        
#         model = Llama(**llm_config)
#         st.success("Model loaded successfully!")
#         return model
#     except Exception as e:
#         st.error(f"Error loading model: {str(e)}")
#         raise

@st.cache_resource(show_spinner=False)
def load_llama_model():
    """Load Llama model with caching"""
    try:
        model_path = "mistral-7b-v0.1.Q4_K_M.gguf"
        
        if not os.path.exists(model_path):
            st.info("Downloading model... This may take a while.")
            direct_url = "https://huggingface.co/TheBloke/Mistral-7B-v0.1-GGUF/resolve/main/mistral-7b-v0.1.Q4_K_M.gguf"
            download_file_with_progress(direct_url, model_path)

        if not os.path.exists(model_path):
            raise FileNotFoundError("Model file not found after download attempt")

        if os.path.getsize(model_path) < 1000000:  # Less than 1MB
            raise ValueError("Model file is too small, likely corrupted")

        llm_config = {
            "model_path": model_path,
            "n_ctx": 2048,
            "n_threads": 4,
            "n_batch": 512,
            "n_gpu_layers": 0,
            "verbose": True  # Enable verbose mode for debugging
        }
        
        logging.info("Initializing Llama model...")
        model = Llama(**llm_config)
        
        # Test the model
        logging.info("Testing model...")
        test_response = model("Test", max_tokens=10)
        if not test_response:
            raise RuntimeError("Model test failed")
            
        logging.info("Model loaded and tested successfully")
        st.success("Model loaded successfully!")
        return model
        
    except Exception as e:
        logging.error(f"Error loading model: {str(e)}")
        logging.error("Full error details: ", exc_info=True)
        raise

def check_environment():
    """Check if the environment is properly set up"""
    try:
        import torch
        import sentence_transformers
        return True
    except ImportError as e:
        st.error(f"Missing required package: {str(e)}")
        st.stop()
        return False

class RAGPipeline:
    def __init__(self, data_folder: str, k: int = 5):
        self.data_folder = data_folder
        self.k = k
        self.retriever = SentenceTransformerRetriever()
        self.documents = []
        self.device = torch.device("cpu")
        self.llm = load_llama_model()

    def preprocess_query(self, query: str) -> str:
        """Clean and prepare the query"""
        query = query.lower().strip()
        query = re.sub(r'\s+', ' ', query)
        return query

    def postprocess_response(self, response: str) -> str:
        """Clean up the generated response"""
        response = response.strip()
        response = re.sub(r'\s+', ' ', response)
        response = re.sub(r'\d{4}-\d{2}-\d{2}\s\d{2}:\d{2}:\d{2}(?:\+\d{2}:?\d{2})?', '', response)
        return response

    def query_model(self, prompt: str) -> str:
        """Query the local Llama model"""
        try:
            if self.llm is None:
                raise RuntimeError("Model not initialized")
                
            response = self.llm(
                prompt,
                max_tokens=512,
                temperature=0.4,
                top_p=0.95,
                echo=False,
                stop=["Question:", "\n\n"]
            )
            
            if response and 'choices' in response and len(response['choices']) > 0:
                text = response['choices'][0].get('text', '').strip()
                return text
            else:
                raise ValueError("No valid response generated")
                
        except Exception as e:
            logging.error(f"Error in query_model: {str(e)}")
            raise

    # def process_query(self, query: str, placeholder) -> str:
    #     try:
    #         # Preprocess query
    #         query = self.preprocess_query(query)
            
    #         # Show retrieval status
    #         status = placeholder.empty()
    #         status.write("🔍 Finding relevant information...")
            
    #         # Get embeddings and search
    #         query_embedding = self.retriever.encode([query])
    #         similarities = F.cosine_similarity(query_embedding, self.retriever.doc_embeddings)
    #         scores, indices = torch.topk(similarities, k=min(self.k, len(self.documents)))
            
    #         relevant_docs = [self.documents[idx] for idx in indices.tolist()]
            
    #         # Update status
    #         status.write("💭 Generating response...")
            
    #         # Prepare context and prompt
    #         context = "\n".join(relevant_docs[:3])
    #         prompt = f"""Context information is below:
    #         {context}
            
    #         Given the context above, please answer the following question:
    #         {query}

    #         Guidelines:
    #         - If you cannot answer based on the context, say so politely
    #         - Keep the response concise and focused
    #         - Only include sports-related information
    #         - No dates or timestamps in the response
    #         - Use clear, natural language
            
    #         Answer:"""
            
    #         # Generate response
    #         response_placeholder = placeholder.empty()
            
    #         try:
    #             response_text = self.query_model(prompt)
    #             if response_text:
    #                 final_response = self.postprocess_response(response_text)
    #                 response_placeholder.markdown(final_response)
    #                 return final_response
    #             else:
    #                 message = "No relevant answer found. Please try rephrasing your question."
    #                 response_placeholder.warning(message)
    #                 return message
                    
    #         except Exception as e:
    #             logging.error(f"Generation error: {str(e)}")
    #             message = "Had some trouble generating the response. Please try again."
    #             response_placeholder.warning(message)
    #             return message
                
    #     except Exception as e:
    #         logging.error(f"Process error: {str(e)}")
    #         message = "Something went wrong. Please try again with a different question."
    #         placeholder.warning(message)
    #         return message

    def process_query(self, query: str, placeholder) -> str:
    try:
        # Preprocess query
        query = self.preprocess_query(query)
        logging.info(f"Processing query: {query}")
        
        # Show retrieval status
        status = placeholder.empty()
        status.write("🔍 Finding relevant information...")
        
        # Get embeddings and search
        query_embedding = self.retriever.encode([query])
        similarities = F.cosine_similarity(query_embedding, self.retriever.doc_embeddings)
        scores, indices = torch.topk(similarities, k=min(self.k, len(self.documents)))
        
        # Log similarity scores
        for idx, score in zip(indices.tolist(), scores.tolist()):
            logging.info(f"Score: {score:.4f} | Document: {self.documents[idx][:100]}...")
        
        relevant_docs = [self.documents[idx] for idx in indices.tolist()]
        
        # Update status
        status.write("💭 Generating response...")
        
        # Prepare context and prompt
        context = "\n".join(relevant_docs[:3])
        prompt = f"""Context information is below:
        {context}
        
        Given the context above, please answer the following question:
        {query}

        Guidelines:
        - If you cannot answer based on the context, say so politely
        - Keep the response concise and focused
        - Only include sports-related information
        - No dates or timestamps in the response
        - Use clear, natural language
        
        Answer:"""
        
        # Generate response
        response_placeholder = placeholder.empty()
        
        try:
            # Add logging for model state
            logging.info("Model state check - Is None?: " + str(self.llm is None))
            
            # Directly use Llama model
            response = self.llm(
                prompt,
                max_tokens=512,
                temperature=0.4,
                top_p=0.95,
                echo=False,
                stop=["Question:", "\n\n"]
            )
            
            logging.info(f"Raw model response: {response}")
            
            if response and isinstance(response, dict) and 'choices' in response:
                generated_text = response['choices'][0].get('text', '').strip()
                if generated_text:
                    final_response = self.postprocess_response(generated_text)
                    response_placeholder.markdown(final_response)
                    return final_response
            
            message = "No relevant answer found. Please try rephrasing your question."
            response_placeholder.warning(message)
            return message
            
        except Exception as e:
            logging.error(f"Generation error: {str(e)}")
            logging.error(f"Full error details: ", exc_info=True)
            message = f"Had some trouble generating the response: {str(e)}"
            response_placeholder.warning(message)
            return message
            
    except Exception as e:
        logging.error(f"Process error: {str(e)}")
        logging.error(f"Full error details: ", exc_info=True)
        message = f"Something went wrong: {str(e)}"
        placeholder.warning(message)
        return message

@st.cache_resource(show_spinner=False)
def initialize_rag_pipeline():
    """Initialize the RAG pipeline once"""
    try:
        # Create necessary directories
        os.makedirs("ESPN_data", exist_ok=True)
        
        # Load embeddings from Drive
        drive_file_id = "1MuV63AE9o6zR9aBvdSDQOUextp71r2NN"
        with st.spinner("Loading embeddings from Google Drive..."):
            cache_data = load_from_drive(drive_file_id)
            if cache_data is None:
                st.error("Failed to load embeddings from Google Drive")
                st.stop()
        
        # Initialize pipeline
        data_folder = "ESPN_data"
        rag = RAGPipeline(data_folder)
        
        # Store embeddings
        rag.documents = cache_data['documents']
        rag.retriever.store_embeddings(cache_data['embeddings'])
        
        return rag
        
    except Exception as e:
        logging.error(f"Pipeline initialization error: {str(e)}")
        st.error(f"Failed to initialize the system: {str(e)}")
        raise

# def main():
#     try:
#         # Environment check
#         if not check_environment():
#             return

#         # Improved CSS styling
#         st.markdown("""
#             <style>
#             /* Container styling */
#             .block-container {
#                 padding-top: 2rem;
#                 padding-bottom: 2rem;
#             }
            
#             /* Text input styling */
#             .stTextInput > div > div > input {
#                 width: 100%;
#             }
            
#             /* Button styling */
#             .stButton > button {
#                 width: 200px;
#                 margin: 0 auto;
#                 display: block;
#                 background-color: #FF4B4B;
#                 color: white;
#                 border-radius: 5px;
#                 padding: 0.5rem 1rem;
#             }
            
#             /* Title styling */
#             .main-title {
#                 text-align: center;
#                 padding: 1rem 0;
#                 font-size: 3rem;
#                 color: #1F1F1F;
#             }
            
#             .sub-title {
#                 text-align: center;
#                 padding: 0.5rem 0;
#                 font-size: 1.5rem;
#                 color: #4F4F4F;
#             }
            
#             /* Description styling */
#             .description {
#                 text-align: center;
#                 color: #666666;
#                 padding: 0.5rem 0;
#                 font-size: 1.1rem;
#                 line-height: 1.6;
#                 margin-bottom: 1rem;
#             }

#             /* Answer container styling */
#             .stMarkdown {
#                 max-width: 100%;
#             }

#             /* Streamlit default overrides */
#             .st-emotion-cache-16idsys p {
#                 font-size: 1.1rem;
#                 line-height: 1.6;
#             }
            
#             /* Container for main content */
#             .main-content {
#                 max-width: 1200px;
#                 margin: 0 auto;
#                 padding: 0 1rem;
#             }
#             </style>
#         """, unsafe_allow_html=True)

#         # Header section
#         st.markdown("<h1 class='main-title'>🏆 The Sport Chatbot</h1>", unsafe_allow_html=True)
#         st.markdown("<h3 class='sub-title'>Using ESPN API</h3>", unsafe_allow_html=True)
#         st.markdown("""
#             <p class='description'>
#                 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.
#             </p>
#             <p class='description'>
#                 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!
#             </p>
#         """, unsafe_allow_html=True)

#         # Initialize the pipeline
#         if 'rag' not in st.session_state:
#             with st.spinner("Loading resources..."):
#                 st.session_state.rag = initialize_rag_pipeline()

#         # Create columns for layout
#         col1, col2, col3 = st.columns([1, 6, 1])
        
#         with col2:
#             # Query input
#             query = st.text_input("What would you like to know about sports?")
            
#             if st.button("Get Answer"):
#                 if query:
#                     response_placeholder = st.empty()
#                     try:
#                         response = st.session_state.rag.process_query(query, response_placeholder)
#                         logging.info(f"Generated response: {response}")
#                     except Exception as e:
#                         logging.error(f"Query processing error: {str(e)}")
#                         response_placeholder.warning("Unable to process your question. Please try again.")
#                 else:
#                     st.warning("Please enter a question!")

#         # Footer
#         st.markdown("<br><br>", unsafe_allow_html=True)
#         st.markdown("---")
#         st.markdown("""
#             <p style='text-align: center; color: #666666; padding: 1rem 0;'>
#                 Powered by ESPN Data & Mistral AI 🚀
#             </p>
#         """, unsafe_allow_html=True)

#     except Exception as e:
#         logging.error(f"Application error: {str(e)}")
#         st.error("An unexpected error occurred. Please check the logs and try again.")

def main():
    try:
        # Environment check
        if not check_environment():
            return

        # Improved CSS styling
        st.markdown("""
            <style>
            /* Container styling */
            .block-container {
                padding-top: 2rem;
                padding-bottom: 2rem;
            }
            
            /* Text input styling */
            .stTextInput > div > div > input {
                width: 100%;
            }
            
            /* Button styling */
            .stButton > button {
                width: 200px;
                margin: 0 auto;
                display: block;
                background-color: #FF4B4B;
                color: white;
                border-radius: 5px;
                padding: 0.5rem 1rem;
            }
            
            /* Title styling */
            .main-title {
                text-align: center;
                padding: 1rem 0;
                font-size: 3rem;
                color: #1F1F1F;
            }
            
            .sub-title {
                text-align: center;
                padding: 0.5rem 0;
                font-size: 1.5rem;
                color: #4F4F4F;
            }
            
            /* Description styling */
            .description {
                text-align: center;
                color: #666666;
                padding: 0.5rem 0;
                font-size: 1.1rem;
                line-height: 1.6;
                margin-bottom: 1rem;
            }

            /* Answer container styling */
            .stMarkdown {
                max-width: 100%;
            }

            /* Streamlit default overrides */
            .st-emotion-cache-16idsys p {
                font-size: 1.1rem;
                line-height: 1.6;
            }
            
            /* Container for main content */
            .main-content {
                max-width: 1200px;
                margin: 0 auto;
                padding: 0 1rem;
            }
            </style>
        """, unsafe_allow_html=True)

        # Header section
        st.markdown("<h1 class='main-title'>🏆 The Sport Chatbot</h1>", unsafe_allow_html=True)
        st.markdown("<h3 class='sub-title'>Using ESPN API</h3>", unsafe_allow_html=True)
        st.markdown("""
            <p class='description'>
                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.
            </p>
            <p class='description'>
                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!
            </p>
        """, unsafe_allow_html=True)

        # Initialize the pipeline with better error handling
        if 'rag' not in st.session_state:
            try:
                with st.spinner("Loading resources..."):
                    st.session_state.rag = initialize_rag_pipeline()
                    logging.info("Pipeline initialized successfully")
            except Exception as e:
                logging.error(f"Pipeline initialization error: {str(e)}")
                st.error("Failed to initialize the system. Please check the logs.")
                st.stop()
                return

        # Create columns for layout
        col1, col2, col3 = st.columns([1, 6, 1])
        
        with col2:
            # Query input
            query = st.text_input("What would you like to know about sports?")
            
            if st.button("Get Answer"):
                if query:
                    response_placeholder = st.empty()
                    try:
                        # Log query processing start
                        logging.info(f"Processing query: {query}")
                        
                        # Process query and get response
                        response = st.session_state.rag.process_query(query, response_placeholder)
                        
                        # Log successful response
                        logging.info(f"Generated response: {response}")
                    except Exception as e:
                        # Log error details
                        logging.error(f"Query processing error: {str(e)}")
                        logging.error("Full error details: ", exc_info=True)
                        response_placeholder.warning("Unable to process your question. Please try again.")
                else:
                    st.warning("Please enter a question!")

        # Footer
        st.markdown("<br><br>", unsafe_allow_html=True)
        st.markdown("---")
        st.markdown("""
            <p style='text-align: center; color: #666666; padding: 1rem 0;'>
                Powered by ESPN Data & Mistral AI 🚀
            </p>
        """, unsafe_allow_html=True)

    except Exception as e:
        logging.error(f"Application error: {str(e)}")
        logging.error("Full error details: ", exc_info=True)
        st.error("An unexpected error occurred. Please check the logs and try again.")

if __name__ == "__main__":
    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    
    try:
        main()
    except Exception as e:
        logging.error(f"Fatal error: {str(e)}")
        logging.error("Full error details: ", exc_info=True)
        st.error("A fatal error occurred. Please check the logs and try again.")

if __name__ == "__main__":
    main()