import os
import json
import argparse
import chromadb
from chromadb.utils import embedding_functions
from mlx_lm import load, generate, stream_generate
from mlx_lm.sample_utils import make_sampler

# ==========================================
# 1. CONFIGURATION SECTION
# ==========================================
# This section defines the paths to our files, models, and databases.
# Having them here makes it easy to change them later if needed.

# Path to the directory where our local Vector Database (ChromaDB) stores its data
CHROMA_DB_PATH = "data/mintoak/chroma_db"

# The name of the collection (table) in our Vector Database
COLLECTION_NAME = "mintoak_content"

# Path to the file containing our prepared text passages (chunks)
CHUNKS_JSON_PATH = "data/mintoak/mintoak_chunks.json"

# The base language model we are loading from Hugging Face / local cache
MODEL_PATH = "mlx-community/Qwen2.5-1.5B-Instruct-4bit"

# Path to the directory containing our trained fine-tuning LoRA adapters
ADAPTER_PATH = "adapters/mintoak"


# ==========================================
# 2. HELPER FUNCTIONS
# ==========================================

def get_embedding_function():
    """
    Creates and returns the embedding model.
    Embeddings convert text sentences into lists of numbers (vectors) 
    so that the computer can search for documents by 'meaning' instead of exact keywords.
    We use the SentenceTransformer 'all-MiniLM-L6-v2' which runs fully locally.
    """
    return embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="all-MiniLM-L6-v2"
    )

def initialize_vector_db():
    """
    Initializes ChromaDB and populates or updates it if content has changed.
    This ensures that when you add new blogs or products to the JSON file, 
    the Vector Database is updated automatically without manual folder deletions.
    """
    # Connect to the local database stored on your disk
    client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    emb_fn = get_embedding_function()
    
    # Get the existing collection, or create a new one if it doesn't exist
    collection = client.get_or_create_collection(
        name=COLLECTION_NAME, 
        embedding_function=emb_fn
    )
    
    # Check if the chunks file exists before trying to read it
    if not os.path.exists(CHUNKS_JSON_PATH):
        print(f"Error: Chunks file '{CHUNKS_JSON_PATH}' not found.")
        print("Please run scripts/mintoak/prepare_data.py first to parse the raw data.")
        return collection

    # Load current chunks from the JSON file
    with open(CHUNKS_JSON_PATH, "r", encoding="utf-8") as f:
        chunks = json.load(f)

    # Programmatically reset and reload the collection if the counts don't match
    # (This triggers automatically when you add new data to your JSON source file)
    if collection.count() != len(chunks):
        print("Detected changes in chunks. Programmatically updating Vector DB...")
        try:
            existing = collection.get()
            if existing and "ids" in existing and existing["ids"]:
                ids_to_del = existing["ids"]
                for j in range(0, len(ids_to_del), 500):
                    collection.delete(ids=ids_to_del[j:j+500])
        except Exception as e:
            print(f"Failed to clear existing database entries: {e}")
        
        # Extract the fields we need to insert
        ids = [c["id"] for c in chunks]
        documents = [f"Title: {c['title']}\nContent: {c['content']}" for c in chunks]
        metadatas = [{"url": c["url"], "title": c["title"], "category": c["category"]} for c in chunks]
        
        # ChromaDB allows batch insertion; we insert in groups of 500
        batch_size = 500
        for i in range(0, len(chunks), batch_size):
            collection.add(
                ids=ids[i:i+batch_size],
                documents=documents[i:i+batch_size],
                metadatas=metadatas[i:i+batch_size]
            )
        print(f"Successfully updated index with {collection.count()} passages.")
    else:
        print(f"Vector DB is up-to-date with {collection.count()} passages.")
        
    return collection


def retrieve_context(collection, query, top_k=3):
    """
    Searches the Vector Database for the most relevant text passages matching the user's question.
    """
    results = collection.query(
        query_texts=[query],
        n_results=top_k
    )
    
    query_lower = query.lower()
    
    # Guardrail: Distance-based Grounding Filter
    # Refuse questions that are semantically unrelated to Mintoak data (high distance)
    # and do not explicitly mention Mintoak or its product lines.
    best_distance = 999.0
    if results and 'distances' in results and len(results['distances'][0]) > 0:
        best_distance = results['distances'][0][0]
        
    in_scope_keywords = [
        # 1. Brand Names
        "mintoak", "digiledge",
        # 2. Product Names
        "smartpayments", "digionboard", "rewardrun", "business360", "sellsmart", "staffaccess", "soundhub", "soundbox",
        # 3. Leadership Names
        "raman", "khanduja", "sanjay", "nazareth", "rama", "tadepalli", "kabeer", "jain", "rohit", "ramana",
        # 4. FinTech Domain
        "payment", "payments", "transaction", "transactions", "merchant", "merchants", "bank", "banks", "card", "cards", 
        "gpv", "volume", "acquiring", "acquirer", "acquirers", "credit", "lending", "loan", "loans", "invoicing", 
        "reconciliation", "settlement", "settlements",
        # 5. MSME/SME
        "sme", "smes", "msme", "msmes", "business", "businesses", "retail", "retailer", "retailers",
        # 6. Onboarding & KYC
        "onboard", "onboarding", "kyc", "kyb", "register", "registration", "account", "accounts", "signup", "signin", 
        "login", "log", "logout",
        # 7. Security & Credentials
        "password", "passwords", "reset", "otp", "credential", "credentials", "auth", "authentication", "security",
        # 8. Contact & Offices
        "address", "addresses", "office", "offices", "contact", "location", "locations", "headquarters", "registered", 
        "corporate", "callback", "call", "support", "help", "demo",
        # 9. Careers
        "job", "jobs", "career", "careers", "hiring", "opening", "openings", "join", "recruit", "recruitment",
        # 10. Reports & Docs
        "report", "reports", "statement", "statements", "download", "downloads", "dashboard",
        # 11. Developer Resources
        "api", "apis", "sdk", "sdks", "developer", "developers"
    ]
    import re
    is_mintoak_mention = any(re.search(rf"\b{re.escape(w)}\b", query_lower) for w in in_scope_keywords)
    
    # Strict cut-off: if distance is extremely high, it's always out-of-scope
    if best_distance > 0.82:
        return "OUT_OF_SCOPE_REFUSAL", {}
    # Moderate cut-off: allow bypass only if it has brand context
    if best_distance > 0.60 and not is_mintoak_mention:
        return "OUT_OF_SCOPE_REFUSAL", {}
    
    # 1. If user is asking about products in general, ensure the master catalog is included
    catalog_doc = None
    catalog_meta = None
    if any(pw in query_lower for pw in ["product", "products", "offering", "offerings", "catalog"]):
        try:
            catalog_res = collection.get(ids=["synthetic_master_catalog"])
            if catalog_res and catalog_res.get('documents'):
                catalog_doc = catalog_res['documents'][0]
                catalog_meta = catalog_res['metadatas'][0]
        except Exception as e:
            print(f"Failed to retrieve product list catalog by ID: {e}")

    context_parts = []
    url_to_title = {}
    
    if catalog_doc and catalog_meta:
        context_parts.append(
            f"Source: {catalog_meta['url']}\n"
            f"Content: {catalog_doc}"
        )
        url_to_title[catalog_meta['url']] = catalog_meta.get('title', 'Mintoak Product List Catalog')
        
    # If we found matches, format each match with its source URL and content
    if results and 'documents' in results and len(results['documents'][0]) > 0:
        indices = list(range(len(results['documents'][0])))
        # Prioritize the Product List Catalog chunk if retrieved
        for i in indices:
            meta = results['metadatas'][0][i]
            if meta.get('title') == "Mintoak Product List Catalog":
                indices.remove(i)
                indices.insert(0, i)
                break
                
        for idx in indices:
            doc = results['documents'][0][idx]
            meta = results['metadatas'][0][idx]
            # Avoid duplicating the catalog if it was already prepended
            if meta.get('title') == "Mintoak Product List Catalog" and catalog_doc is not None:
                continue
            context_parts.append(
                f"Source: {meta['url']}\n"
                f"Content: {doc}"
            )
            url_to_title[meta['url']] = meta.get('title', 'Learn more')
            
    # Join the retrieved passages together separated by dashes
    return "\n\n---\n\n".join(context_parts), url_to_title

# ==========================================
# 3. MAIN EXECUTION FLOW
# ==========================================

def main():
    # Set up command-line arguments so users can pass their query when running the script
    parser = argparse.ArgumentParser(description="Mintoak local RAG Assistant")
    parser.add_argument("query", type=str, help="The question you want to ask the assistant")
    parser.add_argument("--top_k", type=int, default=3, help="Number of content blocks to retrieve")
    args = parser.parse_args()

    # Guardrail Check: Pre-processing Safety & Prompt Injection Filter
    # A lightweight filter to block offensive language and prompt injection attacks before loading the model
    profane_patterns = [
        "fuck", "shit", "asshole", "bitch", "bastard", "cunt", "dick", "pussy",
        "idiot", "stupid", "dumbass", "kill yourself", "kys", "hate speech",
        "retard", "wanker", "motherfucker"
    ]
    injection_patterns = [
        "ignore previous", "system prompt", "developer mode", "you are now", 
        "pretend to be", "reveal instructions", "system instruction", "dan mode", 
        "hypothetically speaking", "override rules", "jailbreak"
    ]
    
    query_lower = args.query.lower()
    
    if any(word in query_lower for word in profane_patterns):
        print("\n--- Assistant Response ---")
        print("I cannot assist with queries containing inappropriate or offensive language.")
        print("--------------------------\n")
        return
        
    if any(pattern in query_lower for pattern in injection_patterns):
        print("\n--- Assistant Response ---")
        print("I am unable to comply with requests that attempt to modify or bypass system instructions.")
        print("--------------------------\n")
        return

    # Step 1: Connect to or update the local Vector Database
    collection = initialize_vector_db()

    # Step 2: Retrieve context if needed, or bypass if it's a general greeting/identity question
    print(f"\nRetrieving relevant content for: '{args.query}'...")
    
    # Check if the question is a basic greeting or identity question
    import re
    query_clean = args.query.strip().lower().rstrip("?").strip()
    
    # Matches variations of greetings like: hi, hii, hello, helloo, hey, heyy, yo, yoo, good morning, etc.
    greeting_regex = r"^(hi+|hello+|hey+|yo+|greetings|good\s+(morning|afternoon|evening))\b"
    is_greeting = bool(re.match(greeting_regex, query_clean))
    
    identity_regex = r"\b(who are (you|u)|what is your name|what's your name|what do you do|introduce yourself|tell me about yourself)\b"
    is_identity = bool(re.search(identity_regex, query_clean))
    
    is_general = is_greeting or is_identity
    
    if is_general:
        print("\n--- Assistant Response ---")
        if is_greeting:
            print("Hello! 👋 I am the Mintoak Website Assistant. How can I help you with Mintoak's products or services today? 😊")
        else:
            print("I am the Mintoak Website Assistant! 🤖 I can help answer questions about Mintoak Innovations, our product catalog, services, office address, and guidelines. How can I assist you today? ✨")
        print("--------------------------\n")
        return
    else:
        # Search the database for relevant blogs/pages matching the user question
        context, url_to_title = retrieve_context(collection, args.query, args.top_k)
        
        if context == "OUT_OF_SCOPE_REFUSAL":
            print("\n--- Assistant Response ---")
            print("I am only authorized to assist with questions related to Mintoak Innovations Private Limited, its products, and services. How can I help you with Mintoak today? 😊")
            print("--------------------------\n")
            return
    
    # The System instruction defines the persona, grounding rule, and formatting rules.
    system_prompt = (
        "You are the Mintoak Website Assistant. Speak in a friendly, fresh, and enthusiastic tone with expressive emojis! "
        "Answer strictly using the provided context. Never invent, assume, or pull from outside knowledge. "
        "Mintoak is strictly a bank-led B2B platform for merchants; Mintoak does NOT issue loyalty points or rewards to end-consumers. If asked how customers or end-consumers earn loyalty points, clarify that Mintoak does not directly run consumer loyalty programs, but enables merchants to set up their own custom campaigns and offers for their customers. Do not confuse B2B merchant app users with B2C end-consumers. "
        "Start your response directly and naturally—do NOT use robotic preambles like 'As per the context' or 'Based on the provided context'. "
        "Be highly concise (maximum 2-3 sentences or a short bulleted list). Cite the source URL from the context at the end of your response. "
        "If the context does not contain the answer, reply: 'The requested information does not currently exist on www.mintoak.com. "
        "You can get in touch with our team at https://www.mintoak.com/contact-us.'"
    )
    
    # Construct the user message matching the exact format the model was trained on
    prompt = f"Context:\n{context}\n\nQuestion: {args.query}"
    
    # Format the messages in standard chat roles
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]
 
    # Step 4: Load the base Qwen model along with our trained LoRA adapters
    print("\nLoading model & adapter (this may take a few seconds on first run)...")
    model, tokenizer = load(MODEL_PATH, adapter_path=ADAPTER_PATH)
 
    # Step 5: Convert chat messages into the model's special token format (Chat Template)
    formatted_prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # Step 6: Stream the response token-by-token for immediate output
    print("\n--- Assistant Response ---")

    import time
    import re
    import sys

    start_time = time.time()
    streamed_tokens = []

    for token_response in stream_generate(
        model,
        tokenizer,
        prompt=formatted_prompt,
        max_tokens=256,
        sampler=make_sampler(temp=0.0)
    ):
        token_text = token_response.text
        streamed_tokens.append(token_text)
        print(token_text, end="", flush=True)

    latency = time.time() - start_time
    print()  # newline after streamed output

    response = "".join(streamed_tokens).strip()

    # Clean up robotic preambles
    preambles = [
        r"^(?:as per|according to|based on)\s+(?:the\s+)?(?:context|provided context|mintoak documentation|documentation)(?:,\s*)?",
        r"^based on what is provided in the context(?:,\s*)?",
        r"^from the context(?:,\s*)?"  
    ]
    for pattern in preambles:
        response = re.compile(pattern, re.IGNORECASE).sub("", response).strip()

    if response and response[0].islower():
        response = response[0].upper() + response[1:]

    # Clean up learned dataset prefixes from the response
    prefix_pattern = r"^Based on the Mintoak documentation,\s*(?:here is what we know about\s+[^:\n]+:)?\s*"
    response = re.sub(prefix_pattern, "", response, flags=re.IGNORECASE).strip()

    # If the response doesn't include a Source citation, append markdown source links
    suffix = ""
    if "Source:" not in response and context:
        sources = []
        for line in context.split("\n"):
            if line.startswith("Source: "):
                url = line.replace("Source: ", "").strip()
                if url not in sources:
                    sources.append(url)
        if sources:
            markdown_sources = []
            for url in sources:
                title = url_to_title.get(url, "Learn more")
                markdown_sources.append(f"[{title}]({url})")
            suffix = "\n\n👉 For more details, visit: " + ", ".join(markdown_sources)

    # Print only the suffix (source links) since the response was already streamed live above
    if suffix:
        print(suffix)

    print(f"\n⏱️  Response generated in {latency:.2f} seconds.")
    print("--------------------------\n")


if __name__ == "__main__":
    main()