import os
import re
import json
import random

SCRAPED_DIR = r"c:\Risu Solutions\ByteAstra\backend\data\ayurveda\scraped"
OUTPUT_FILE = r"c:\Risu Solutions\ByteAstra\backend\data\finetune\ayurveda_train_real.jsonl"

def parse_md_file(filepath):
    """
    Parses a markdown file to extract frontmatter and body.
    """
    with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
        content = f.read()

    # Match frontmatter
    fm_match = re.match(r"^---\s*\n(.*?)\n---\s*\n", content, re.DOTALL)
    if not fm_match:
        return None, content

    fm_text = fm_match.group(1)
    body = content[fm_match.end():]

    # Parse YAML-like fields
    metadata = {}
    for line in fm_text.split("\n"):
        if ":" in line:
            k, v = line.split(":", 1)
            metadata[k.strip().lower()] = v.strip()

    return metadata, body

def chunk_text(text, min_words=100, max_words=300):
    """
    Splits text by double newlines, merging short paragraphs to form stable chunks.
    """
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    chunks = []
    current_chunk = []
    current_word_count = 0

    for p in paragraphs:
        word_count = len(p.split())
        if current_word_count + word_count <= max_words:
            current_chunk.append(p)
            current_word_count += word_count
        else:
            if current_chunk and current_word_count >= min_words:
                chunks.append("\n\n".join(current_chunk))
                current_chunk = [p]
                current_word_count = word_count
            elif not current_chunk:
                # Paragraph itself is larger than max_words, add it directly
                chunks.append(p)
            else:
                chunks.append("\n\n".join(current_chunk))
                current_chunk = [p]
                current_word_count = word_count

    if current_chunk:
        chunks.append("\n\n".join(current_chunk))

    return chunks

def extract_topic(chunk):
    """
    Extracts a topic name from the chunk (first heading or bolded word, or first few words).
    """
    lines = chunk.split("\n")
    # Check for short heading-like lines
    for line in lines[:3]:
        line_clean = line.strip()
        if line_clean and len(line_clean) < 60 and not line_clean.startswith("-") and not re.match(r"^\d", line_clean):
            return line_clean.replace("**", "").replace("##", "").strip()
            
    # Check for bolded word
    bold_match = re.search(r"\*\*(.*?)\*\*", chunk)
    if bold_match:
        return bold_match.group(1)

    # Fallback to first 5 words
    words = chunk.split()
    if len(words) > 5:
        return " ".join(words[:5]) + "..."
    return "Ayurvedic Concepts"

def build_dataset():
    random.seed(42)
    entries = []

    print("Scanning scraped Ayurveda files...")
    for root, dirs, files in os.walk(SCRAPED_DIR):
        for file in files:
            if file.endswith(".md"):
                filepath = os.path.join(root, file)
                metadata, body = parse_md_file(filepath)
                if not metadata:
                    continue

                source = metadata.get("source", "Classical Ayurveda Texts")
                section = metadata.get("section", "General Chapters")
                chapter = metadata.get("chapter", "Core Syllabus")

                chunks = chunk_text(body)
                for chunk in chunks:
                    word_count = len(chunk.split())
                    if word_count < 40:
                        continue # Skip extremely short items/noise
                        
                    topic = extract_topic(chunk)
                    
                    # Generate structural queries
                    queries = [
                        f"Show me the verbatim text from {source}, {section}, {chapter} concerning '{topic}'.",
                        f"Provide the classical verses/text describing '{topic}' in {source} ({chapter}, {section}).",
                        f"What does {source} say in {chapter} ({section}) about '{topic}'?",
                        f"Give the textual reference for '{topic}' from {source}, {section}, {chapter}.",
                        f"Explain the topic of '{topic}' as documented in {source} ({section}, {chapter})."
                    ]
                    query = random.choice(queries)

                    system_prompt = "You are ByteAstra, a precise BAMS tutor. Answer questions using only the provided classical textbooks context. Quote the textbook exactly when requested."
                    
                    entry = {
                        "conversations": [
                            {"role": "system", "content": system_prompt},
                            {"role": "user", "content": query},
                            {"role": "assistant", "content": chunk}
                        ]
                    }
                    entries.append(entry)

    # Shuffle the dataset
    random.shuffle(entries)
    
    # Save JSONL
    os.makedirs(os.path.dirname(OUTPUT_FILE), exist_ok=True)
    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        for entry in entries:
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")

    print(f"Dataset compiled successfully! Generated {len(entries)} training samples in {OUTPUT_FILE}.")

if __name__ == "__main__":
    build_dataset()