import json
import os
import random
import shutil

# ==========================================
# 1. PATH CONFIGURATIONS
# ==========================================
# Paths to input source files and output generated datasets

# Path where the raw blogs/products JSON file is stored
RAW_DATA_PATH = "data/mintoak/mintoak_data.json"

# Path where the training set JSONL will be saved (used for fine-tuning)
TRAIN_OUT_PATH = "data/mintoak/train.jsonl"

# Path where the validation set JSONL will be saved (used to evaluate fine-tuning)
VALID_OUT_PATH = "data/mintoak/valid.jsonl"

# Path where the clean chunked passages are saved (used by RAG vector database)
RAG_OUT_PATH = "data/mintoak/mintoak_chunks.json"


# ==========================================
# 2. DATA PROCESSING FUNCTIONS
# ==========================================

def load_raw_data(path):
    """
    Loads and normalizes raw website content.
    Supports both standard simplified format (title, url, content) 
    and complex nested Strapi CMS export structures.
    """
    if not os.path.exists(path):
        print(f"Error: Raw data file {path} not found.")
        print("Please place your Mintoak products/blogs JSON file at that path.")
        return None
        
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)
        
    # Check if the data is wrapped in a Strapi CMS 'data' dict or is a direct list
    if isinstance(data, dict) and "data" in data:
        raw_items = data["data"]
    elif isinstance(data, list):
        raw_items = data
    else:
        raw_items = []
        
    normalized_items = []
    for item in raw_items:
        if not isinstance(item, dict):
            continue
            
        # Format 1: If it's already in the Simplified Format, load it directly
        if "url" in item or "content" in item:
            normalized_items.append({
                "url": item.get("url", "https://www.mintoak.com"),
                "title": item.get("title", "Mintoak Page"),
                "content": item.get("content", ""),
                "category": item.get("category", "general")
            })
            continue
            
        # Format 2: Parse nested Strapi CRM exporter format
        attributes = item.get("attributes") or {}
        
        # Format 2A: Product format
        product_content_list = attributes.get("product") or []
        if product_content_list:
            slug = attributes.get("slug") or ""
            url = f"https://www.mintoak.com/products/{slug}" if slug else "https://www.mintoak.com"
            title = attributes.get("breadcrumbTitle")
            content_parts = []
            category = "Product offering"
            
            for p_item in product_content_list:
                if not isinstance(p_item, dict):
                    continue
                # Extract description or listing text
                desc = p_item.get("description") or p_item.get("ListingShortDescription") or p_item.get("metaDescription")
                if desc:
                    content_parts.append(desc)
                # Try to extract title if not set yet
                if not title and p_item.get("__component") == "shared.seo":
                    title = p_item.get("metaTitle")
            
            if not title:
                title = "Mintoak Product"
                
            content = "\n\n".join(content_parts)
            
            normalized_items.append({
                "url": url.strip(),
                "title": title.strip(),
                "content": content.strip(),
                "category": category
            })
            continue

        # Format 2C: About Us / Visionaries format
        about_us_list = attributes.get("aboutUs") or []
        if about_us_list:
            url = "https://www.mintoak.com/about-us"
            title = None
            for au_item in about_us_list:
                if not isinstance(au_item, dict):
                    continue
                sec_titles = au_item.get("sectionTitle") or []
                for st in sec_titles:
                    if isinstance(st, dict) and st.get("title"):
                        title = st.get("title")
                        break
                if title:
                    break
            if not title:
                title = "Mintoak Founders"
                
            content_parts = []
            for au_item in about_us_list:
                if not isinstance(au_item, dict):
                    continue
                if au_item.get("__component") == "about-us.mintoak-visionaries":
                    team_members = au_item.get("teamMembers") or []
                    members_descriptions = []
                    for member in team_members:
                        name = member.get("title", "").strip()
                        role = member.get("subTitle", "").strip()
                        desc = member.get("description", "").strip()
                        members_descriptions.append(f"{name} ({role} - {desc})")
                    
                    content_parts.append("Mintoak founders: " + ", ".join(members_descriptions) + ".")
            
            content = " ".join(content_parts)
            category = "Company Info"
            
            normalized_items.append({
                "url": url.strip(),
                "title": title.strip(),
                "content": content.strip(),
                "category": category
            })
            continue

        # Format 2B: Blog format
        slug = attributes.get("slug") or ""
        url = f"https://www.mintoak.com/blog/{slug}" if slug else "https://www.mintoak.com"
        
        title = None
        # Extract title from components list
        top_title_desc = attributes.get("TopTittleDescription") or []
        for title_desc in top_title_desc:
            if isinstance(title_desc, dict) and title_desc.get("__component") == "reference.title":
                title = title_desc.get("title")
                break
        if not title:
            # Fallback to SEO meta title if description title is missing
            seo = attributes.get("seo") or {}
            title = seo.get("metaTitle", "Mintoak Page")
            
        # Get category label safely
        category_obj = attributes.get("blog_master_category") or {}
        category_data = category_obj.get("data") or {}
        category_attr = category_data.get("attributes") or {}
        category = category_attr.get("category", "general")
        
        # Merge all paragraphs under BlogContent into a single block of text
        content_parts = []
        blog_content = attributes.get("BlogContent") or []
        for content_item in blog_content:
            if isinstance(content_item, dict):
                description = content_item.get("description")
                if description:
                    content_parts.append(description)
        content = "\n\n".join(content_parts)
        
        normalized_items.append({
            "url": url.strip(),
            "title": title.strip() if title else "Mintoak Page",
            "content": content.strip(),
            "category": category.strip() if category else "general"
        })
        
    return normalized_items


def chunk_content(text, max_chars=800):
    """
    Splits long articles into smaller passages (chunks) of roughly max_chars length.
    This ensures that when a passage is retrieved, it fits within the model's memory, 
    and keeps sentences intact so sentences don't get chopped in half.
    """
    sentences = text.split(". ")
    chunks = []
    current_chunk = []
    current_len = 0
    for s in sentences:
        s_len = len(s)
        # If adding this sentence exceeds chunk size limit, save current chunk and start a new one
        if current_len + s_len > max_chars:
            if current_chunk:
                chunks.append(". ".join(current_chunk) + ".")
            current_chunk = [s]
            current_len = s_len
        else:
            current_chunk.append(s)
            current_len += s_len + 2
    if current_chunk:
        chunks.append(". ".join(current_chunk) + ".")
    return chunks


def build_rag_chunks(raw_data):
    """
    Iterates through all normalized articles and splits them into chunks with metadata.
    """
    chunks_db = []
    for idx, item in enumerate(raw_data):
        url = item.get("url", "https://www.mintoak.com")
        title = item.get("title", "Mintoak Page")
        content = item.get("content", "")
        category = item.get("category", "general")
        
        # Split content into smaller passages
        chunk_size = 1200 if category == "Company Info" else 800
        passages = chunk_content(content, max_chars=chunk_size)
        for p_idx, passage in enumerate(passages):
            chunks_db.append({
                "id": f"doc_{idx}_chunk_{p_idx}",
                "url": url,
                "title": title,
                "content": passage.strip(),
                "category": category
            })
    return chunks_db


def generate_synthetic_dialogue(chunks_db):
    """
    Generates synthetic training pairs to teach the model:
    1. Grounding: Answer strictly based on the provided context.
    2. Citations: Reference the exact URL/section.
    3. Fallback: Refuse to answer if not in the context.
    4. General/Greetings: Correctly answer greetings and identity questions.
    5. Lead Capture: Model conversational ask-then-tag flow.
    """
    qa_pairs = []
    
    # Prompt template for fine-tuning system instructions
    system_prompt = (
        "You are the Mintoak Website Assistant, a senior, business-aware digital advisor on the Mintoak website. "
        "Your tone is professional, clear, confident, technology-forward, and benefit-led. Do not use hype, casualness, or over-enthusiasm. "
        "Do not use emojis by default; use at most one emoji only where it adds clarity (e.g. 💡, 📈) and never in pricing, credibility, or trust contexts. "
        "Answer strictly using the provided context. Never invent, assume, or pull from outside knowledge. "
        "Start your response directly and naturally—do NOT use robotic preambles like 'As per the context' or 'Based on the provided context'. "
        "Frame Mintoak as a white-labeled, modular, cloud-native, API-first SaaS platform/solution provider (never as a 'software company', 'software', 'tool', or 'app'). "
        "Avoid these banned phrases: 'seamless', 'empowering', 'leverage' (as a verb), 'game-changer', 'synergy', 'unlocking', 'in today\'s fast-paced world', 'as we move forward'. "
        "Follow this exact structure: Context (direct answer) -> Insight (translate features into merchant/business outcomes like activation, cost-to-acquire, time-to-market) -> Direction (nudge to the next best page/demo and close with the source URL). "
        "Keep responses highly concise (maximum 2-3 sentences or short bullets) unless explicitly asked for detail. "
        "Output tail order: answer -> source URL(s) -> [CAPTURE_LEAD] (if applicable). "
        "If the context does not contain the answer, reply exactly: '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.'"
    )

    def clean_banned_words(text):
        import re
        # Case-insensitive replacement of banned words
        text = re.sub(r'\bseamlessly\b', 'efficiently', text, flags=re.I)
        text = re.sub(r'\bseamless\b', 'integrated', text, flags=re.I)
        text = re.sub(r'\bempowers\b', 'enables', text, flags=re.I)
        text = re.sub(r'\bempowering\b', 'enabling', text, flags=re.I)
        text = re.sub(r'\bempower\b', 'enable', text, flags=re.I)
        text = re.sub(r'\bleveraged\b', 'utilized', text, flags=re.I)
        text = re.sub(r'\bleverages\b', 'utilizes', text, flags=re.I)
        text = re.sub(r'\bleveraging\b', 'utilizing', text, flags=re.I)
        text = re.sub(r'\bleverage\b', 'use', text, flags=re.I)
        text = re.sub(r'\bgame-changer\b', 'significant advancement', text, flags=re.I)
        text = re.sub(r'\bgame changer\b', 'significant advancement', text, flags=re.I)
        text = re.sub(r'\bsynergy\b', 'alignment', text, flags=re.I)
        text = re.sub(r'\bsynergies\b', 'alignments', text, flags=re.I)
        text = re.sub(r'\bunlocking\b', 'releasing', text, flags=re.I)
        text = re.sub(r'\bunlock\b', 'access', text, flags=re.I)
        text = re.sub(r"\bin today's fast-paced world,?\s*", "", text, flags=re.I)
        text = re.sub(r"\bas we move forward,?\s*", "", text, flags=re.I)
        
        text = re.sub(r'\bMintoak (\w+ )?(?:software|tool|app)\b', r'Mintoak \1platform', text, flags=re.I)
        text = re.sub(r'\b(?:software|tool|app) Mintoak\b', 'Mintoak platform', text, flags=re.I)
        
        # Force Mintoak capitalization
        text = re.sub(r'(?<![\./])\b(MintOak|MINTOAK|mintoak)\b(?![\./])', 'Mintoak', text)
        return text

    # 1. Generate Positive Examples (grounded QA)
    for chunk in chunks_db:
        context_str = f"Source: {chunk['url']}\nContent: {chunk['content']}"
        
        # Filter rules to target content matching Mintoak or merchant topics
        title_lower = chunk['title'].lower()
        content_lower = chunk['content'].lower()
        category_lower = chunk['category'].lower()
        
        if "360" in title_lower or "merchant" in content_lower or "product" in category_lower or "qr" in title_lower:
            questions = [
                f"What is {chunk['title']}?",
                f"Can you explain {chunk['title']} and its benefits?",
                f"How does {chunk['title']} help merchants?",
                f"Tell me about {chunk['title']}."
            ]
            q = random.choice(questions)
            
            # Extract first 2 sentences
            sentences = [s.strip() for s in chunk['content'].split('.') if s.strip()]
            summary_sentences = sentences[:2]
            summary_text = ". ".join(summary_sentences)
            if summary_text and not summary_text.endswith('.'):
                summary_text += '.'
                
            # Structured Response: Context -> Insight -> Direction
            context_ans = clean_banned_words(summary_text)
            
            insight_choices = [
                "This enables banks to deep-dive into merchant needs and support business growth.",
                "This module helps acquirers improve merchant retention and increase transaction metrics.",
                "By optimizing these services, banks can reduce cost-to-acquire and accelerate merchant time-to-market.",
                "This serves as a key driver for business growth, helping merchants scale their daily operations."
            ]
            insight_text = random.choice(insight_choices)
            
            direction_text = f"👉 For more details, visit: [{chunk['title']}]({chunk['url']})"
            
            a = f"{context_ans}\n\n{insight_text}\n\n{direction_text}"
            
            qa_pairs.append({
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"Context:\n{context_str}\n\nQuestion: {q}"},
                    {"role": "assistant", "content": a}
                ]
            })
            
        # 1B. Generate Catalog List Examples (for list/summary queries)
        if chunk['id'] == "synthetic_master_catalog":
            catalog_questions = [
                "give me list of mintoak product offerings",
                "give me list of mintoak product offerings and explain each one liner",
                "what are the product offerings of mintoak",
                "list all products of mintoak and their descriptions",
                "give me all product names and one liners",
                "what products does mintoak offer"
            ]
            for cq in catalog_questions:
                catalog_cleaned = clean_banned_words(chunk['content'])
                ca = f"{catalog_cleaned}\n\n👉 For more details, visit: [{chunk['title']}]({chunk['url']})"
                qa_pairs.append({
                    "messages": [
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": f"Context:\n{context_str}\n\nQuestion: {cq}"},
                        {"role": "assistant", "content": ca}
                    ]
                })

    # 2. Generate Negative/Fallback Examples (out-of-domain QA)
    out_of_domain_questions = [
        "What is the weather in Mumbai today?",
        "How do I sign up for a Netflix account?",
        "Tell me about the payments page or FAQs.",
        "What is the capital of France?",
        "Can I get a refund on my purchase?"
    ]
    
    fallback_response = (
        "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."
    )

    for q in out_of_domain_questions:
        unrelated_context = "Source: https://www.mintoak.com\nContent: Mintoak is a SaaS platform partnering with banks to support merchants."
        qa_pairs.append({
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Context:\n{unrelated_context}\n\nQuestion: {q}"},
                {"role": "assistant", "content": fallback_response}
            ]
        })

    # 3. Generate General/Greeting/Identity Examples (positive handling)
    general_examples = [
        {
            "q": "who are you?",
            "a": "I am the Mintoak Website Assistant, a digital advisor here to help you understand our SaaS platform and merchant solutions."
        },
        {
            "q": "who are u?",
            "a": "I am the Mintoak Website Assistant. How can I help you today?"
        },
        {
            "q": "what is your name?",
            "a": "I am the Mintoak Website Assistant."
        },
        {
            "q": "hi",
            "a": "Hello! I am the Mintoak Website Assistant. How can I help you today?"
        },
        {
            "q": "hello",
            "a": "Hello! I am the Mintoak Website Assistant. How can I help you today?"
        },
        {
            "q": "what is mintoak?",
            "a": "Mintoak is a white-labeled, modular, cloud-native SaaS platform that partners with banks and acquirers to enable and support merchants. It provides integrated modules for digital payment acceptance, reconciliation, settlements, and business analytics."
        },
        {
            "q": "what is mintoak's offering?",
            "a": "Mintoak offers a white-labeled, modular SaaS platform for banks and acquirers. Key modules include digital payment acceptance, automated transaction reconciliation, settlement tracking, and data-driven analytics to support merchants."
        }
    ]

    for item in general_examples:
        qa_pairs.append({
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Context:\n\nQuestion: {item['q']}"},
                {"role": "assistant", "content": item['a']}
            ]
        })

    # 4. Generate Conversational Lead Capture Examples
    lead_capture_examples = [
        {
            "history": [
                {"role": "user", "content": "I would like to schedule a demo of Mintoak DigiOnboard"}
            ],
            "a": "We would be glad to show you how DigiOnboard supports merchant acquisition. Could you please share your name and email address so our business team can set up a live demonstration?"
        },
        {
            "history": [
                {"role": "user", "content": "How much does Mintoak SmartPayments cost?"}
            ],
            "a": "Pricing for SmartPayments is customized based on your transaction volume and deployment needs. Could you please provide your name and email address so we can share a detailed proposal?"
        },
        {
            "history": [
                {"role": "user", "content": "How much does Engage360 cost?"},
                {"role": "assistant", "content": "Pricing for Engage360 is tailored to each acquirer's deployment. Could you please share your name and email address so we can provide a customized quote?"},
                {"role": "user", "content": "Sure, my name is Alice Smith and my email is alice@bank.com"}
            ],
            "a": "Thank you, Alice. A Mintoak representative will reach out to you shortly at alice@bank.com with our custom pricing structure. For more details, visit: [Mintoak Engage360](https://www.mintoak.com/products/mintoak-engage360) [CAPTURE_LEAD]"
        },
        {
            "history": [
                {"role": "user", "content": "Can you contact me about partnering? My name is Bob and email is bob@acquirer.com"}
            ],
            "a": "Thank you for your interest in partnering with Mintoak, Bob. Our partnerships team will contact you shortly at bob@acquirer.com to discuss collaboration opportunities. You can learn more about us at: [About Us](https://www.mintoak.com/about-us) [CAPTURE_LEAD]"
        }
    ]

    for item in lead_capture_examples:
        messages = [{"role": "system", "content": system_prompt}]
        for hist_item in item["history"]:
            messages.append({"role": hist_item["role"], "content": hist_item["content"]})
        qa_pairs.append({
            "messages": messages + [{"role": "assistant", "content": item["a"]}]
        })

    return qa_pairs


# ==========================================
# 3. MAIN RUNNER
# ==========================================

def main():
    # Load and clean the raw data
    raw_data = load_raw_data(RAW_DATA_PATH)
    if not raw_data:
        return
        
    print(f"Loaded {len(raw_data)} website items.")

    # Dynamically extract and combine all office addresses under contact-us
    office_contents = []
    other_items = []
    contact_us_item = None
    
    for item in raw_data:
        url = item.get("url", "")
        if url.startswith("https://www.mintoak.com/contact-us#"):
            office_contents.append(item.get("content", "").strip())
        elif url == "https://www.mintoak.com/contact-us":
            contact_us_item = item
        else:
            other_items.append(item)
            
    if office_contents:
        combined_address_content = "Mintoak office locations and addresses:\n\n" + "\n\n".join(f"* {content}" for content in office_contents)
        if contact_us_item:
            contact_us_item["content"] = combined_address_content
            contact_us_item["title"] = "Mintoak Office Addresses"
        else:
            contact_us_item = {
                "url": "https://www.mintoak.com/contact-us",
                "title": "Mintoak Office Addresses",
                "content": combined_address_content,
                "category": "Company Info"
            }
            
    if contact_us_item:
        other_items.append(contact_us_item)
        
    raw_data = other_items
    
    # Process text into chunks
    chunks_db = build_rag_chunks(raw_data)
    
    # Programmatically compile Master Catalog Chunk with one-liner descriptions (deduplicated)
    product_dict = {}
    for item in raw_data:
        if item.get('category') == 'Product offering':
            p_title = item.get('title', '').strip()
            p_content = item.get('content', '').strip()
            
            # Extract the first sentence or first line as a clean one-liner
            first_sentence = p_content.split('\n')[0].split('. ')[0].strip()
            if first_sentence and not first_sentence.endswith('.'):
                first_sentence += '.'
                
            # Limit description length for the catalog summary chunk
            if len(first_sentence) > 160:
                first_sentence = first_sentence[:157] + "..."
                
            if p_title and first_sentence and p_title not in product_dict:
                product_dict[p_title] = first_sentence
                
    product_details = [f"* {k}: {v}" for k, v in sorted(product_dict.items())]
                
    if product_details:
        catalog_content = "Mintoak offers the following product offerings and descriptions:\n" + "\n".join(product_details)
        chunks_db.append({
            "id": "synthetic_master_catalog",
            "url": "https://www.mintoak.com/products",
            "title": "Mintoak Product List Catalog",
            "content": catalog_content,
            "category": "Product offering"
        })
        print(f"Generated synthetic master catalog chunk with {len(product_details)} unique product one-liners.")
        
    print(f"Generated {len(chunks_db)} RAG chunks.")
    
    # Save the chunks to the RAG database source file
    with open(RAG_OUT_PATH, 'w', encoding='utf-8') as f:
        json.dump(chunks_db, f, indent=4)
    print(f"Saved RAG index database to {RAG_OUT_PATH}")
    
    # Clear Vector DB cache to force rebuild on next assistant run
    shutil.rmtree("data/mintoak/chroma_db", ignore_errors=True)
    print("Cleared Vector DB cache directory to force index rebuild.")
    
    # Generate dialogue pairs for training/validation
    qa_pairs = generate_synthetic_dialogue(chunks_db)
    print(f"Generated {len(qa_pairs)} training/validation pairs.")
    
    # Shuffle and split into 80% train and 20% validation sets
    random.shuffle(qa_pairs)
    split_idx = int(len(qa_pairs) * 0.8)
    train_data = qa_pairs[:split_idx]
    valid_data = qa_pairs[split_idx:]
    
    # Save datasets in JSONL (JSON Lines) format required by MLX fine-tuning
    with open(TRAIN_OUT_PATH, 'w', encoding='utf-8') as f:
        for item in train_data:
            f.write(json.dumps(item) + "\n")
            
    with open(VALID_OUT_PATH, 'w', encoding='utf-8') as f:
        for item in valid_data:
            f.write(json.dumps(item) + "\n")
            
    print(f"Saved {len(train_data)} training examples to {TRAIN_OUT_PATH}")
    print(f"Saved {len(valid_data)} validation examples to {VALID_OUT_PATH}")

if __name__ == "__main__":
    main()