"""
ByteAstra — 100-Question BAMS Stress Test
Difficulty: Easy (25) | Medium (25) | Hard (25) | Very Tough (25)
Topics: Doshas, Dhatus, Malas, Agni, Panchakarma, Prakriti, Srotas,
        Panchamahabhutas, Dinacharya, Rasas, Ama, Ojas, Samprapti, Nidana
"""
import asyncio
import time
import json
import sys

sys.path.insert(0, "/code")
from app.services import engine

QUESTIONS = [
    # ── EASY (1-25) ──────────────────────────────────────────────────────────
    ("easy",  1, "What is Ayurveda?"),
    ("easy",  2, "What are the three Doshas?"),
    ("easy",  3, "What is Vata dosha?"),
    ("easy",  4, "What is Pitta dosha?"),
    ("easy",  5, "What is Kapha dosha?"),
    ("easy",  6, "What are the seven Dhatus?"),
    ("easy",  7, "What are the three Malas?"),
    ("easy",  8, "What is Agni in Ayurveda?"),
    ("easy",  9, "What is Ama?"),
    ("easy", 10, "What is Ojas?"),
    ("easy", 11, "What is Prakriti?"),
    ("easy", 12, "What are Panchamahabhutas?"),
    ("easy", 13, "What is Dinacharya?"),
    ("easy", 14, "What is Panchakarma?"),
    ("easy", 15, "What is Jatharagni?"),
    ("easy", 16, "What is Rasa Dhatu?"),
    ("easy", 17, "What is Rakta Dhatu?"),
    ("easy", 18, "What is Mamsa Dhatu?"),
    ("easy", 19, "What is Meda Dhatu?"),
    ("easy", 20, "What is Asthi Dhatu?"),
    ("easy", 21, "What is Majja Dhatu?"),
    ("easy", 22, "What is Shukra Dhatu?"),
    ("easy", 23, "What are the six Rasas (tastes) in Ayurveda?"),
    ("easy", 24, "What is Vamana therapy?"),
    ("easy", 25, "What is Virechana therapy?"),

    # ── MEDIUM (26-50) ───────────────────────────────────────────────────────
    ("medium", 26, "List the five subtypes of Vata with their locations and functions."),
    ("medium", 27, "List the five subtypes of Pitta with their locations and functions."),
    ("medium", 28, "List the five subtypes of Kapha with their locations and functions."),
    ("medium", 29, "Explain the 13 types of Agni with their roles in digestion."),
    ("medium", 30, "What are the qualities (Gunas) of Vata dosha?"),
    ("medium", 31, "What are the qualities (Gunas) of Pitta dosha?"),
    ("medium", 32, "What are the qualities (Gunas) of Kapha dosha?"),
    ("medium", 33, "Describe the concept of Srotas and list the major ones."),
    ("medium", 34, "What is Vikriti and how does it differ from Prakriti?"),
    ("medium", 35, "Explain the relationship between Agni and Ama formation."),
    ("medium", 36, "Describe Sama Agni and its significance."),
    ("medium", 37, "What is Vishama Agni and what dosha causes it?"),
    ("medium", 38, "What is Tikshna Agni and what dosha causes it?"),
    ("medium", 39, "What is Manda Agni and what dosha causes it?"),
    ("medium", 40, "Describe the properties of Akasha (space element)."),
    ("medium", 41, "Describe the properties of Vayu (air element)."),
    ("medium", 42, "Describe the properties of Teja (fire element)."),
    ("medium", 43, "Describe the properties of Jala (water element)."),
    ("medium", 44, "Describe the properties of Prithvi (earth element)."),
    ("medium", 45, "What is Basti therapy and what dosha does it primarily treat?"),
    ("medium", 46, "What is Nasya therapy and what dosha does it primarily treat?"),
    ("medium", 47, "What is Raktamokshana and when is it indicated?"),
    ("medium", 48, "Explain the concept of Ahara Rasa and how it forms Rasa Dhatu."),
    ("medium", 49, "What is Vata Prakriti and what are its physical and mental characteristics?"),
    ("medium", 50, "What is Pitta Prakriti and what are its characteristics?"),

    # ── HARD (51-75) ─────────────────────────────────────────────────────────
    ("hard", 51, "Explain the Dhatu-Agni pathway: how Rasa Dhatu transforms into Shukra Dhatu step by step."),
    ("hard", 52, "What are the 5 Bhutagnis and how do they relate to the Panchamahabhutas?"),
    ("hard", 53, "Describe Samprapti (pathogenesis) with all six stages of disease development."),
    ("hard", 54, "Compare and contrast Vata, Pitta and Kapha Prakriti in terms of body build, skin, digestion, sleep and mental nature."),
    ("hard", 55, "Explain the functions of Prana Vata in detail."),
    ("hard", 56, "Explain the functions of Udana Vata in detail."),
    ("hard", 57, "Explain the functions of Samana Vata in detail."),
    ("hard", 58, "Explain the functions of Apana Vata in detail."),
    ("hard", 59, "Explain the functions of Vyana Vata in detail."),
    ("hard", 60, "Describe Alochaka Pitta with location and function."),
    ("hard", 61, "Describe Sadhaka Pitta with location and function."),
    ("hard", 62, "Describe Pachaka Pitta with location and function."),
    ("hard", 63, "Describe Ranjaka Pitta with location and function."),
    ("hard", 64, "Describe Bhrajaka Pitta with location and function."),
    ("hard", 65, "Explain the role of Ojas in immunity and its relationship to all seven Dhatus."),
    ("hard", 66, "What is the Mahakoshtha and how does it relate to Vata disorders?"),
    ("hard", 67, "Describe the Ashtavidha Pariksha (eightfold examination) in detail."),
    ("hard", 68, "Explain how Ama differs from Aamavastha and what are its clinical features."),
    ("hard", 69, "Describe the Purvakarma (preparatory procedures) before Panchakarma."),
    ("hard", 70, "What are the contraindications of Vamana therapy?"),
    ("hard", 71, "Describe the Nidana Panchaka and its role in diagnosis."),
    ("hard", 72, "What is Dhatvagni and how does each Dhatvagni produce Sara, Mala and Upadhatu?"),
    ("hard", 73, "Explain the concept of Satva, Rajas and Tamas Manasa Prakriti."),
    ("hard", 74, "Describe the 20 Gurvadi Gunas (20 qualities) and their therapeutic significance."),
    ("hard", 75, "What is Agni theory as described in Charaka Samhita Chapter 14?"),

    # ── VERY TOUGH (76-100) ──────────────────────────────────────────────────
    ("very_tough", 76, "Explain the concept of Tridosha Siddhanta and how the interaction of three Doshas maintains homeostasis."),
    ("very_tough", 77, "Describe the Shad-Kriya-Kala (six stages of Samprapti) with examples for Prameha."),
    ("very_tough", 78, "What are Dushyas and how do they interact with Doshas in Dosha-Dushya Sammurchana?"),
    ("very_tough", 79, "Describe the complete pathogenesis (Samprapti) of Jwara (fever) including all stages."),
    ("very_tough", 80, "Explain Trividha Pariksha (threefold examination) and its clinical application."),
    ("very_tough", 81, "Describe all seven Dhatus with their Agni, Sara (essence), Mala (waste) and Upadhatu (byproduct)."),
    ("very_tough", 82, "What is the significance of Kala (time) in Ayurvedic treatment planning?"),
    ("very_tough", 83, "Explain the theory of Loka-Purusha Samya (macrocosm-microcosm relationship)."),
    ("very_tough", 84, "Describe the Srotas, their Mula (root), Marga (path) and Mukha (opening) for all 13 Srotamsi."),
    ("very_tough", 85, "What are the three types of Ama according to Dosha dominance and their symptoms?"),
    ("very_tough", 86, "Explain the role of Vata in the nervous system and compare it with modern neurophysiology."),
    ("very_tough", 87, "Describe Rasayana therapy: definition, types, indications and examples from classical texts."),
    ("very_tough", 88, "What is Vajikarana? Describe its types and classical formulations."),
    ("very_tough", 89, "Explain the Panchabhouta composition of each Dosha and how it determines their qualities."),
    ("very_tough", 90, "Describe the complete process of Snehakarma (oleation therapy) with types, methods and indications."),
    ("very_tough", 91, "What is Shodhana Chikitsa and how does it differ from Shamana Chikitsa?"),
    ("very_tough", 92, "Describe the Chaturvidha Pariksha Bhava (four diagnostic tools) according to Charaka."),
    ("very_tough", 93, "Explain the role of Pitta in metabolism and compare with modern hepatic function."),
    ("very_tough", 94, "What is the Anuvasana Basti and how does it differ from Niruha Basti?"),
    ("very_tough", 95, "Describe all types of Virechana dravyas with their Dosha specificity."),
    ("very_tough", 96, "Explain Kapha Prakriti characteristics and diseases of Kapha dominance."),
    ("very_tough", 97, "What is the significance of Agni in Ayurvedic oncology (Arbuda Chikitsa)?"),
    ("very_tough", 98, "Describe the complete Dinacharya routine with timings and therapeutic rationale."),
    ("very_tough", 99, "Explain the concept of Satmya and how it applies to diet and lifestyle modification."),
    ("very_tough", 100, "What is Swasthavritta? Describe all its components including Dinacharya, Ritucharya and Sadvritta."),
]

RESULTS = []

async def run_one(level, qnum, query):
    start = time.time()
    first_tok = None
    words = 0
    full = []
    grounded = False
    citations = 0
    error = None

    try:
        async for event in engine.stream_answer("ayurveda", query, []):
            t = event.get("type")
            if t == "citations":
                grounded = event.get("is_grounded", False)
                citations = len(event.get("citations", []))
            elif t == "delta":
                c = event.get("content", "")
                if c:
                    if first_tok is None:
                        first_tok = time.time() - start
                    words += len(c.split())
                    full.append(c)
            elif t == "error":
                error = event.get("error")
    except Exception as e:
        error = str(e)

    total = time.time() - start
    response = "".join(full).strip()

    # Quality heuristics
    has_heading = "##" in response or "###" in response
    has_list = "- " in response or "1." in response
    has_sanskrit = any(w in response for w in ["Agni","Dosha","Dhatu","Vata","Pitta","Kapha","Ama","Ojas"])
    score_hint = "GOOD" if (words > 80 and has_heading and has_list and has_sanskrit) else ("PARTIAL" if words > 40 else "SHORT")

    result = {
        "level": level,
        "q": qnum,
        "query": query,
        "grounded": grounded,
        "citations": citations,
        "ttft": round(first_tok, 1) if first_tok else None,
        "total_s": round(total, 1),
        "words": words,
        "quality": score_hint,
        "error": error,
        "response_preview": response[:200] + "..." if len(response) > 200 else response,
    }
    RESULTS.append(result)

    print(f"[Q{qnum:03d}|{level.upper()[:4]}] {total:.0f}s | {words}w | {score_hint} | Grounded={grounded} | {query[:60]}")
    return result

async def main():
    print("="*70)
    print("ByteAstra 100-Question BAMS Stress Test")
    print("="*70)

    for level, qnum, query in QUESTIONS:
        await run_one(level, qnum, query)
        await asyncio.sleep(2)  # brief pause between questions

    # Summary
    print("\n" + "="*70)
    print("STRESS TEST SUMMARY")
    print("="*70)

    by_level = {}
    for r in RESULTS:
        by_level.setdefault(r["level"], []).append(r)

    for level in ["easy", "medium", "hard", "very_tough"]:
        group = by_level.get(level, [])
        grounded_pct = sum(1 for r in group if r["grounded"]) / max(len(group), 1) * 100
        good_pct = sum(1 for r in group if r["quality"] == "GOOD") / max(len(group), 1) * 100
        avg_time = sum(r["total_s"] for r in group) / max(len(group), 1)
        avg_words = sum(r["words"] for r in group) / max(len(group), 1)
        errors = sum(1 for r in group if r["error"])
        print(f"\n{level.upper():<12} | n={len(group)} | Grounded={grounded_pct:.0f}% | Good={good_pct:.0f}% | AvgTime={avg_time:.0f}s | AvgWords={avg_words:.0f} | Errors={errors}")

    total_grounded = sum(1 for r in RESULTS if r["grounded"])
    total_good = sum(1 for r in RESULTS if r["quality"] == "GOOD")
    total_errors = sum(1 for r in RESULTS if r["error"])
    avg_t = sum(r["total_s"] for r in RESULTS) / max(len(RESULTS), 1)
    print(f"\nOVERALL | n={len(RESULTS)} | Grounded={total_grounded/len(RESULTS)*100:.0f}% | Good={total_good/len(RESULTS)*100:.0f}% | AvgTime={avg_t:.0f}s | Errors={total_errors}")

    # Save JSON
    with open("/code/stress_results.json", "w") as f:
        json.dump(RESULTS, f, indent=2)
    print("\nFull results saved to /code/stress_results.json")

if __name__ == "__main__":
    asyncio.run(main())