"""
Gradio app for Hugging Face chatbot with RAG capabilities.
"""
import warnings

# Suppress deprecation from dependencies (e.g. accelerate) until they use torch.distributed.ReduceOp
warnings.filterwarnings(
    "ignore",
    message=".*torch.distributed.reduce_op.*ReduceOp.*",
    category=FutureWarning,
)

import re
import gradio as gr
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes
import os
from typing import List, Tuple
from huggingface_hub import InferenceClient
from ingestion import DocumentIngestion


# Create a clean minimalist theme
class MinimalistTheme(Base):
    """A clean, minimalist theme with subtle colors and simple styling."""
    def __init__(self):
        super().__init__(
            primary_hue=colors.blue,
            secondary_hue=colors.gray,
            neutral_hue=colors.gray,
            spacing_size=sizes.spacing_md,
            radius_size=sizes.radius_sm,
            text_size=sizes.text_md,
            font=(
                fonts.GoogleFont("Inter"),
                "ui-sans-serif",
                "system-ui",
                "sans-serif",
            ),
            font_mono=(
                fonts.GoogleFont("JetBrains Mono"),
                "ui-monospace",
                "monospace",
            ),
        )
        super().set(
            # Clean backgrounds
            background_fill_primary="#ffffff",
            background_fill_primary_dark="#1a1a1a",
            background_fill_secondary="#f3f6ee",
            background_fill_secondary_dark="#24372b",
            body_background_fill="#f3f6ee",
            body_background_fill_dark="#ffffff",
            block_background_fill="#f3f6ee",
            block_background_fill_dark="#24372b",
            
            # Subtle borders
            block_border_width="1px",
            block_border_color="#e0e0e0",
            block_border_color_dark="#2a2a2a",
            block_shadow="none",
            
            # Clean buttons
            button_primary_background_fill="#8abf50",
            button_primary_background_fill_hover="#1d4ed8",
            button_primary_text_color="#ffffff",
            button_primary_background_fill_dark="#8abf50",
            button_primary_background_fill_hover_dark="#2563eb",
            button_secondary_background_fill="#247b55",
            button_secondary_background_fill_hover="#e5e7eb",
            button_secondary_text_color="#ffffff",
            button_secondary_background_fill_dark="#374151",
            button_secondary_background_fill_hover_dark="#4b5563",
            button_border_width="0px",

            # Input fields
            input_background_fill="#ffffff",
            input_background_fill_dark="#ffffff",
            input_border_width="1px",
            input_border_color="#d1d5db",
            input_border_color_dark="#374151",

            # Text colors
            body_text_color="#1e1e1e",
            body_text_color_dark="#e5e7eb",
            block_label_text_color="#1e1e1e",
            block_label_text_color_dark="#9ca3af",
        )


class RAGChatbot:
    """Chatbot with RAG capabilities."""
    
    # Default and fallback models (try in order until one is supported by your Inference API providers)
    DEFAULT_CHAT_MODEL = "HuggingFaceH4/zephyr-7b-beta"
    FALLBACK_CHAT_MODELS = [
        "mistralai/Mixtral-8x7B-Instruct-v0.1",
        "meta-llama/Llama-3.2-3B-Instruct",
        "Qwen/Qwen2.5-7B-Instruct",
    ]

    def __init__(
        self,
        model_name: str = None,
        embedding_model: str = "all-mpnet-base-v2",
        vector_store_path: str = "data/vector_store"
    ):
        """
        Initialize the RAG chatbot.
        
        Args:
            model_name: Hugging Face model name for the chatbot (via Inference API)
            embedding_model: Model for document embeddings
            vector_store_path: Path to saved vector store
        """
        self.model_name = model_name if model_name else self.DEFAULT_CHAT_MODEL
        # Build list of models to try (primary first, then fallbacks not already primary)
        self._models_to_try = [self.model_name] + [
            m for m in self.FALLBACK_CHAT_MODELS if m != self.model_name
        ]

        # Initialize Inference API client (no model in constructor so we can try multiple)
        hf_token = os.environ.get("HF_TOKEN")
        # Debug: report HF_TOKEN status (masked)
        if not hf_token:
            print("[DEBUG] HF_TOKEN: not set (empty or missing)")
            print("Warning: HF_TOKEN not set. Inference API calls may fail.")
            print("Set HF_TOKEN environment variable or add it to Space secrets.")
        else:
            masked = f"{hf_token[:4]}...{hf_token[-4:]}" if len(hf_token) > 8 else "****"
            print(f"[DEBUG] HF_TOKEN: set (length={len(hf_token)}, masked={masked})")
            print("HF_TOKEN found. Inference API ready.")
        
        print(f"[DEBUG] Inference API client (models to try: {self._models_to_try})")
        try:
            self.inference_client = InferenceClient(token=hf_token)
            print("[DEBUG] Inference API client initialized (model chosen per request with fallbacks)")
        except Exception as e:
            print(f"[DEBUG] Error initializing Inference API client: {type(e).__name__}: {e}")
            self.inference_client = None
        
        # Initialize document ingestion
        self.ingestion = DocumentIngestion(embedding_model=embedding_model)
        
        # Load vector store if it exists
        if os.path.exists(vector_store_path) and os.path.exists(
            os.path.join(vector_store_path, "index.faiss")
        ):
            try:
                self.ingestion.load(vector_store_path)
                print("Loaded existing vector store")
            except Exception as e:
                print(f"Could not load vector store: {e}")
        
        self.chat_history = []
    
    def _generate_with_chat(self, user_content: str, max_new_tokens: int = 512) -> str:
        """Call the Inference API using chat_completion; try fallback models if current is not supported."""
        last_error = None
        for model in self._models_to_try:
            print(f"[DEBUG] _generate_with_chat: trying model={model}, prompt_len={len(user_content)}, max_tokens={max_new_tokens}")
            try:
                response = self.inference_client.chat_completion(
                    model=model,
                    messages=[{"role": "user", "content": user_content}],
                    max_tokens=max_new_tokens,
                    temperature=0.7,
                )
                print(f"[DEBUG] chat_completion OK for model={model}, response type: {type(response).__name__}")
                if response and response.choices and len(response.choices) > 0:
                    msg = response.choices[0].message
                    if hasattr(msg, "content") and msg.content:
                        # Remember this model for next time
                        self.model_name = model
                        self._models_to_try = [model] + [m for m in self._models_to_try if m != model]
                        return msg.content.strip()
                print("[DEBUG] chat_completion returned empty or unexpected structure")
            except Exception as e:
                last_error = e
                err_str = str(e).lower()
                if (
                    "model_not_supported" in err_str
                    or "not supported by any provider" in err_str
                    or "410" in err_str
                    or "gone" in err_str
                ):
                    print(f"[DEBUG] Model {model} not available, trying next fallback.")
                    continue
                print(f"[DEBUG] _generate_with_chat exception for {model}: {type(e).__name__}: {e}")
                import traceback
                traceback.print_exc()
                raise
        if last_error is not None:
            raise last_error
        return ""
    
    def generate_response(self, query: str, use_rag: bool = True, num_results: int = 5) -> str:
        """
        Generate a response to the user query using RAG and Inference API.
        
        Args:
            query: User's question
            use_rag: Whether to use RAG (retrieve relevant documents)
            num_results: Number of document chunks to retrieve
            
        Returns:
            Generated response
        """
        if self.inference_client is None:
            return "Error: Inference API client not initialized. Please check HF_TOKEN configuration."
        
        # If RAG is enabled and we have a vector store, retrieve context and generate answer
        if use_rag and self.ingestion.index is not None:
            try:
                results = self.ingestion.search(query, k=num_results)
                if results:
                    # Build context from retrieved chunks; include source/title so the model can cite it
                    context_parts = []
                    for i, result in enumerate(results, 1):
                        text = result['text'].strip()
                        if not text:
                            continue
                        meta = result.get('metadata') or {}
                        source_label = meta.get('document_title') or meta.get('source') or f"Source {i}"
                        context_parts.append(f"[Context {i}] (Source: {source_label})\n{text}")
                    
                    context = "\n\n".join(context_parts)
                    
                    # Build instruction-tuned prompt
                    prompt = f"""
*You are an expert assistant specializing in organic farming, in particular in Canada and its legal context.
Answer the user's question using only the information provided in the context.
If the context does not include the information needed to answer the question, clearly say:
"The provided context does not contain enough information to answer this question."
Do not alter or paraphrase this exact phrase.
When answering:

Respond in English only.
Do not use outside knowledge, assumptions, or guesswork.
Do not reference or name the source documents anywhere in your answer.
Provide concise, accurate, and helpful explanations.
Do not reveal your internal reasoning. Provide only the final answer.

Structure your answer in the following format:
Summary — A brief, high‑level answer.
Supporting Details — Explain using information only from the provided context. Do not cite or name sources inline.

Context:
{context}

Question: {query}

Answer:"""
                    
                    # Build mapping from context index to source label (and URL if applicable)
                    context_index_to_source = {}
                    context_index_to_url = {}
                    for i, result in enumerate(results, 1):
                        meta = result.get("metadata") or {}
                        context_index_to_source[i] = (
                            meta.get("document_title") or meta.get("source") or f"Source {i}"
                        )
                        if meta.get("type") == "url" and meta.get("source"):
                            context_index_to_url[i] = meta["source"]
                        elif meta.get("url"):
                            context_index_to_url[i] = meta["url"]

                    # Generate response using chat/comversational API (Mistral instruct uses this)
                    try:
                        response_text = self._generate_with_chat(prompt, max_new_tokens=512)
                        if response_text:
                            # Strip all inline [Context N] references and bare "Context N" mentions from the body
                            response_text = re.sub(r'\[Context\s+\d+\]', '', response_text)
                            response_text = re.sub(r'(?<!\[)\bContext\s+[\d][,\s\d]*', '', response_text)
                            response_text = re.sub(r'\s{2,}', ' ', response_text)
                            # Only append references when the model actually answered
                            no_answer_phrase = "The provided context does not contain enough information to answer this question."
                            no_answer_replacement = "The organic documents provided do not contain enough information to answer your question. For more information, try taking a look at the [Canadian Food Inspection Agency](https://inspection.canada.ca/en/food-labels/organic-products/), the [Certified Organic Growers of Canada](https://cog.ca/), or the [Canada Organic Trade Association](https://canada-organic.ca/en/organic)."
                            if no_answer_phrase in response_text:
                                response_text = response_text.replace(no_answer_phrase, no_answer_replacement)
                            if no_answer_replacement not in response_text:
                                # Deduplicate: one entry per unique source URL (or label)
                                seen = {}
                                for i, source_label in context_index_to_source.items():
                                    url = context_index_to_url.get(i)
                                    key = url or source_label
                                    if key not in seen:
                                        seen[key] = (source_label, url)
                                ref_lines = [
                                    "",
                                    "---",
                                    "**References**",
                                ]
                                for idx, (key, (source_label, url)) in enumerate(seen.items(), 1):
                                    if url:
                                        ref_lines.append(f"{idx}. [{source_label}]({url})")
                                    else:
                                        ref_lines.append(f"{idx}. {source_label}")
                                response_text = response_text.rstrip() + "\n\n" + "\n".join(ref_lines)
                            return response_text
                        raise ValueError("Empty response from model")
                    except Exception as api_error:
                        print(f"[DEBUG] RAG generation failed: {type(api_error).__name__}: {api_error}")
                        err_str = str(api_error).lower()
                        if "model_not_supported" in err_str or "not supported by any provider" in err_str:
                            return (
                                "None of the configured chat models are available with your Inference API providers.\n\n"
                                "**How to fix:**\n"
                                "1. See which models are available: https://huggingface.co/inference/models\n"
                                "2. Enable providers (and pick a chat model): https://huggingface.co/settings/inference-api\n"
                                "3. In app.py, set RAGChatbot(model_name=\"your-chosen-model-id\") to match a model you enabled."
                            )
                        # Fallback: return formatted chunks with note
                        response_parts = []
                        response_parts.append("I retrieved relevant information, but couldn't generate a synthesized answer. Here are the relevant chunks:\n\n")
                        for i, result in enumerate(results, 1):
                            meta = result.get('metadata') or {}
                            source = meta.get('document_title') or meta.get('source', '')
                            text = result['text'].strip()
                            if text:
                                response_parts.append(f"**Relevant information {i}** (from {source}):\n{text}\n")
                        return "\n".join(response_parts)
                else:
                    # No results found
                    return "I couldn't find any relevant information in the documents to answer your question. Please try rephrasing or check if the documents contain information about this topic."
            except Exception as e:
                print(f"Error in RAG retrieval: {e}")
                return f"I encountered an error while searching the documents: {str(e)}"
        
        # If no RAG or no vector store, generate response without context
        try:
            prompt = f"""You are a helpful assistant. Answer the following question concisely.

Question: {query}

Answer:"""
            response_text = self._generate_with_chat(prompt, max_new_tokens=256)
            if response_text:
                return response_text
            return "I couldn't generate a response. Please try again."
        except Exception as e:
            print(f"Error generating response: {e}")
            return f"I encountered an error while generating a response: {str(e)}. Please check your HF_TOKEN configuration."
    
    def chat(self, message: str, history):
        """
        Handle chat interaction.
        
        Args:
            message: User message
            history: Chat history (list of ChatMessage or dicts with 'role' and 'content')
            
        Returns:
            Updated history
        """
        if not message or not message.strip():
            return "", history or []
        
        # Ensure history is a list
        if history is None:
            history = []
        
        # Add user message as dictionary
        history.append({"role": "user", "content": message})
        
        # Generate response (always use RAG)
        try:
            response = self.generate_response(message, use_rag=True)
            # Ensure response is not empty
            if not response or not response.strip():
                response = "I'm sorry, I couldn't generate a response. Please try again."
        except Exception as e:
            print(f"Error generating response: {e}")
            import traceback
            traceback.print_exc()
            response = f"I encountered an error: {str(e)}"
        
        # Add assistant response as dictionary
        history.append({"role": "assistant", "content": response})
        
        print(f"Debug - History length: {len(history)}")
        print(f"Debug - Response: {response[:100] if response else 'None'}...")
        
        return "", history


# Initialize chatbot
chatbot = RAGChatbot()


# Create Gradio interface
custom_css = """
.gradio-container textarea {
    color: #1e1e1e !important;
}
"""

with gr.Blocks(title="Organic Certification Assistant", css=custom_css) as app:
    gr.Markdown("<h1 style='color: #24372b; font-size: 2.5rem;'>👤 Organic Certification Assistant</h1>")


    chatbot_interface = gr.Chatbot(
        label="Chat",
        height=500,
        value=[{"role": "assistant", "content": "Welcome to the Organic Certification Assistant! Ask me any questions you have about organic certification and operation in Canada."}]
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="Your Message",
            placeholder="Ask a question about Canadian organics...",
            scale=4
        )
    
    with gr.Row():
        submit_btn = gr.Button("Send", variant="primary")
        clear_btn = gr.Button("Clear")

    with gr.Accordion("Example questions", open=False):
        gr.Markdown("""
- What are the general principles of organic production in Canada?
- What substances are permitted for use in organic crop production?
- Can I use synthetic pesticides on an organic farm?
- What are the requirements for transitioning land to organic certification?
- What livestock practices are required under Canadian organic standards?
- Are antibiotics allowed in organic livestock production?
- What labelling requirements apply to organic products in Canada?
- What is the difference between "organic" and "made with organic ingredients" on a label?
- What are the permitted substances for organic aquaculture in Canada?
- Who certifies organic products in Canada?

> **Disclaimer:** AI-generated responses may not always be accurate or complete. Always verify the information provided against the original source documents and consult official resources before making decisions.
""")
    
    msg.submit(
        chatbot.chat,
        inputs=[msg, chatbot_interface],
        outputs=[msg, chatbot_interface]
    )
    submit_btn.click(
        chatbot.chat,
        inputs=[msg, chatbot_interface],
        outputs=[msg, chatbot_interface]
    )
    def clear_chat():
        return [{"role": "assistant", "content": "Welcome to the Organic Certification Assistant! Ask me any questions you have about organic certification and operation in Canada."}], ""
    
    clear_btn.click(clear_chat, outputs=[chatbot_interface, msg])


if __name__ == "__main__":
    # Get port from environment variable (Hugging Face Spaces sets this) or default to 7860
    port = int(os.environ.get("PORT", 7860))
    app.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=port,
        theme=MinimalistTheme()
    )