"""
ByteAstra — QA Pair Generator for Fine-Tuning Dataset

Reads scraped Ayurveda markdown files and generates instruction-tuning
question-answer pairs in ShareGPT format.

Two modes:
  1. --template-mode  : Rule-based QA extraction (fast, no GPU needed)
  2. --llm-mode       : Uses a local LLM to generate diverse QA pairs (slower, richer)

Usage:
    # Fast template mode (recommended for initial dataset):
    python scripts/generate_qa_pairs.py \
        --input-dir ./data/ayurveda/scraped \
        --output ./data/finetune/ayurveda_train.jsonl \
        --template-mode

Output:
    JSONL file where each line is a JSON conversation object:
    {
      "conversations": [
        {"role": "system",    "content": "..."},
        {"role": "user",      "content": "...question..."},
        {"role": "assistant", "content": "...answer...\n\nSources: Charaka Samhita, Ch. 1"}
      ],
      "metadata": {"source": "...", "chunk_id": "..."}
    }
"""
from __future__ import annotations
import argparse
import hashlib
import json
import logging
import re
import sys
from pathlib import Path

import yaml

logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")
logger = logging.getLogger(__name__)

SYSTEM_PROMPT = (
    "You are ByteAstra, an expert BAMS (Bachelor of Ayurvedic Medicine and Surgery) tutor. "
    "You answer questions about Ayurveda based strictly on classical texts and BAMS curriculum. "
    "Always cite the source textbook and chapter at the end of your answer. "
    "If you do not know the answer from the provided context, say so clearly."
)

# ── Template-based QA patterns ─────────────────────────────────────────────────

QUESTION_TEMPLATES = [
    # Definition questions
    ("What is {concept}?", "definition"),
    ("Define {concept} in Ayurveda.", "definition"),
    ("What does the term '{concept}' mean in Ayurvedic medicine?", "definition"),

    # Explain questions
    ("Explain {concept} according to Ayurveda.", "explanation"),
    ("Describe the concept of {concept} as per classical texts.", "explanation"),
    ("How is {concept} described in Ayurvedic texts?", "explanation"),

    # Function/significance questions
    ("What is the significance of {concept} in maintaining health?", "significance"),
    ("What are the functions of {concept} in the body?", "significance"),
    ("How does {concept} contribute to overall health?", "significance"),

    # List questions
    ("What are the types of {concept}?", "list"),
    ("List the subtypes of {concept} with their characteristics.", "list"),
    ("How many types of {concept} are described in Ayurveda?", "list"),

    # Clinical questions
    ("What are the signs of {concept} imbalance?", "clinical"),
    ("How does disturbance in {concept} manifest as disease?", "clinical"),
    ("What happens when {concept} is aggravated or depleted?", "clinical"),

    # Relationship questions
    ("How is {concept} related to the Doshas?", "relationship"),
    ("What is the relationship between {concept} and digestion?", "relationship"),
    ("How do {concept} and Agni interact in Ayurveda?", "relationship"),
]

# Key Ayurvedic concepts to extract and generate questions around
AYURVEDA_CONCEPTS = [
    "Vata", "Pitta", "Kapha", "Tridosha", "Dosha",
    "Agni", "Jatharagni", "Dhatvagni", "Bhutagni",
    "Rasa Dhatu", "Rakta Dhatu", "Mamsa Dhatu", "Meda Dhatu",
    "Asthi Dhatu", "Majja Dhatu", "Shukra Dhatu", "Sapta Dhatu",
    "Ojas", "Ama", "Srotas", "Prakriti", "Vikriti",
    "Panchamahabhutas", "Akasha", "Vayu", "Teja", "Jala", "Prithvi",
    "Dinacharya", "Ritucharya", "Ahara", "Vihara",
    "Nidana", "Samprapti", "Chikitsa", "Aushadha",
    "Vata Prakriti", "Pitta Prakriti", "Kapha Prakriti",
    "Sama Agni", "Vishama Agni", "Tikshna Agni", "Manda Agni",
]


def parse_frontmatter(text: str) -> tuple[dict, str]:
    meta: dict = {}
    if text.startswith("---"):
        lines = text.splitlines()
        end_idx = None
        for i, line in enumerate(lines[1:], 1):
            if line.strip() == "---":
                end_idx = i
                break
        if end_idx:
            try:
                meta = yaml.safe_load("\n".join(lines[1:end_idx])) or {}
            except Exception:
                pass
            text = "\n".join(lines[end_idx + 1:]).strip()
    return meta, text


def extract_relevant_section(body: str, concept: str, window: int = 600) -> str | None:
    """Find the most relevant paragraph(s) mentioning the concept."""
    concept_lower = concept.lower()
    paragraphs = [p.strip() for p in re.split(r"\n{2,}", body) if len(p.strip()) > 50]

    scored = []
    for para in paragraphs:
        score = para.lower().count(concept_lower)
        if score > 0:
            scored.append((score, para))

    if not scored:
        return None

    # Pick the top paragraph(s) up to window characters
    scored.sort(reverse=True)
    result = []
    total = 0
    for _, para in scored[:3]:
        if total + len(para) <= window:
            result.append(para)
            total += len(para)

    return "\n\n".join(result) if result else None


def build_qa_pair(
    question: str,
    answer_context: str,
    meta: dict,
    chunk_id: str,
) -> dict:
    source = meta.get("source", "Ayurvedic Classical Texts")
    chapter = meta.get("chapter", "")
    url = meta.get("url", "")

    citation = f"\n\nSources: {source}"
    if chapter:
        citation += f", {chapter}"

    full_answer = answer_context.strip() + citation

    return {
        "conversations": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": question},
            {"role": "assistant", "content": full_answer},
        ],
        "metadata": {
            "source": source,
            "chapter": chapter,
            "chunk_id": chunk_id,
            "url": url,
        },
    }


def generate_pairs_for_file(file_path: Path) -> list[dict]:
    text = file_path.read_text(encoding="utf-8", errors="ignore")
    meta, body = parse_frontmatter(text)

    pairs = []
    seen_questions: set[str] = set()

    for concept in AYURVEDA_CONCEPTS:
        context = extract_relevant_section(body, concept)
        if not context or len(context) < 100:
            continue

        # Generate up to 3 questions per concept per file
        import random
        templates = random.sample(QUESTION_TEMPLATES, min(3, len(QUESTION_TEMPLATES)))
        for template, _ in templates:
            question = template.format(concept=concept)
            if question in seen_questions:
                continue
            seen_questions.add(question)

            chunk_id = hashlib.md5(f"{file_path.name}:{concept}:{question}".encode()).hexdigest()[:12]
            pair = build_qa_pair(question, context, meta, chunk_id)
            pairs.append(pair)

    return pairs


def generate_pairs_from_chapters(body: str, meta: dict, file_path: Path) -> list[dict]:
    """Also generate 'full chapter' QA pairs — overview questions about the whole chapter."""
    chapter = meta.get("chapter", file_path.stem)
    source = meta.get("source", "Ayurvedic Classical Texts")

    # Pick the first 800 chars as context for overview question
    overview_context = body[:800].strip()
    if len(overview_context) < 100:
        return []

    overview_questions = [
        f"What is covered in the chapter '{chapter}' of {source}?",
        f"Summarize the key concepts discussed in {chapter}.",
        f"What are the main teachings of {chapter} according to {source}?",
    ]

    pairs = []
    for q in overview_questions[:2]:
        chunk_id = hashlib.md5(f"{file_path.name}:overview:{q}".encode()).hexdigest()[:12]
        pair = build_qa_pair(q, overview_context, meta, chunk_id)
        pairs.append(pair)

    return pairs


def main():
    parser = argparse.ArgumentParser(description="Generate QA pairs from scraped Ayurveda texts")
    parser.add_argument("--input-dir", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--template-mode", action="store_true", default=True,
                        help="Use rule-based QA generation (no GPU needed)")
    parser.add_argument("--max-files", type=int, default=None, help="Limit number of files")
    args = parser.parse_args()

    files = list(args.input_dir.rglob("**/*.md")) + list(args.input_dir.rglob("**/*.txt"))
    if args.max_files:
        files = files[:args.max_files]

    if not files:
        logger.error("No .md/.txt files found in %s", args.input_dir)
        sys.exit(1)

    logger.info("Processing %d files...", len(files))
    args.output.parent.mkdir(parents=True, exist_ok=True)

    all_pairs: list[dict] = []
    for i, file_path in enumerate(files, 1):
        logger.info("[%d/%d] %s", i, len(files), file_path.name)
        text = file_path.read_text(encoding="utf-8", errors="ignore")
        meta, body = parse_frontmatter(text)

        pairs = generate_pairs_for_file(file_path)
        pairs += generate_pairs_from_chapters(body, meta, file_path)
        all_pairs.extend(pairs)
        logger.info("  → %d QA pairs", len(pairs))

    # Deduplicate by question
    seen: set[str] = set()
    unique_pairs = []
    for pair in all_pairs:
        q = pair["conversations"][1]["content"]
        if q not in seen:
            seen.add(q)
            unique_pairs.append(pair)

    # Write JSONL
    with open(args.output, "w", encoding="utf-8") as f:
        for pair in unique_pairs:
            f.write(json.dumps(pair, ensure_ascii=False) + "\n")

    logger.info("✓ Done! %d unique QA pairs written to %s", len(unique_pairs), args.output)
    logger.info("Next step: python finetune/train.py --dataset %s", args.output)


if __name__ == "__main__":
    main()