# scripts/run_audit_diagnostics.py — Comprehensive Audit & Verification Script
import os
import sys
import json
import time
from pathlib import Path

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from knowledge.loader import KnowledgeLoader
from knowledge.chunker import MarkdownChunker
from knowledge.index import HybridIndex
from knowledge.retrieval import HybridRetriever
from knowledge.context import ContextBuilder
from routing.query_router import QueryRouter
from routing.provider_router import ProviderRouter
from policies.university_rules import PolicyEngine, COORDINATION_DISCLAIMER
from policies.language import LanguageEngine
from team.developers import DeveloperRegistry
from cache.manager import CacheManager
from cache.keys import HNU_KNOWLEDGE_VERSION

BASE_DIR = Path(__file__).resolve().parent.parent
DATA_DIR = BASE_DIR / "university_data_arabic"


def audit_section_2_inventory():
    print("\n==================================================")
    print("2. VERIFY DOCUMENT COVERAGE (INVENTORY)")
    print("==================================================")
    
    loader = KnowledgeLoader(data_dir=str(DATA_DIR))
    docs = loader.load_documents()
    summary = loader.get_summary()

    md_files = list(DATA_DIR.rglob("*.md"))
    txt_files = list(DATA_DIR.rglob("*.txt"))
    other_files = [p for p in DATA_DIR.rglob("*") if p.is_file() and p.suffix not in [".md", ".txt"]]

    print(f"MD files found: {len(md_files)}")
    print(f"TXT files found: {len(txt_files)}")
    print(f"Other supported/ignored files: {len(other_files)}")
    print(f"Total source files: {len(md_files) + len(txt_files) + len(other_files)}")
    print(f"Total indexed: {summary['indexed_count']}")
    print(f"Total skipped: {summary['skipped_count']}")
    print(f"Reasons for skipped files: {summary['skipped_files']}")

    print("\nExact indexed filenames list:")
    indexed_names = sorted([d.metadata.filename for d in docs])
    for idx, name in enumerate(indexed_names, 1):
        print(f"  {idx:02d}. {name}")

    data_dir_files = list((DATA_DIR / "data").glob("*"))
    print(f"\nuniversity_data_arabic/data/ files count: {len(data_dir_files)}")
    assert "extracted_table.md" in indexed_names, "extracted_table.md MUST be indexed!"
    assert "website_developers.md" in indexed_names, "website_developers.md MUST be indexed!"
    print("✅ Full recursive document coverage verified!")
    return docs, indexed_names


def audit_section_3_chunking(docs):
    print("\n==================================================")
    print("3. VERIFY CHUNK COVERAGE")
    print("==================================================")

    chunker = MarkdownChunker()
    chunks = chunker.chunk_all(docs)

    chunk_sizes = [len(c.text) for c in chunks]
    avg_size = sum(chunk_sizes) / len(chunk_sizes) if chunk_sizes else 0
    min_size = min(chunk_sizes) if chunk_sizes else 0
    max_size = max(chunk_sizes) if chunk_sizes else 0

    print(f"Total documents: {len(docs)}")
    print(f"Total chunks: {len(chunks)}")
    print(f"Average chunk size: {round(avg_size, 1)} chars")
    print(f"Minimum chunk size: {min_size} chars")
    print(f"Maximum chunk size: {max_size} chars")

    print("\nRandom Chunk Samples:")
    samples = [chunks[0], chunks[len(chunks)//2], chunks[-1]]
    for idx, sample in enumerate(samples, 1):
        print(f"\n--- Sample {idx} ---")
        print(f"Document: {sample.filename}")
        print(f"Section: {sample.section}")
        print(f"Heading Context: {sample.heading_context}")
        print(f"Chunk Text Preview:\n{sample.text[:200]}...")

    return chunks


def audit_section_4_multi_fact(retriever):
    print("\n==================================================")
    print("4. TEST MULTI-FACT RETRIEVAL")
    print("==================================================")

    test_queries = [
        ("A. single fact", "مصاريف كلية طب الأسنان كام؟"),
        ("B. two facts from same doc", "مصاريف طب الأسنان كام وهل ينفع تقسيط؟"),
        ("C. two facts from diff docs", "ما هي مصاريف الطب البشري وما هي الأوراق المطلوبة للتقديم؟"),
        ("D. list questions", "ما هي البرامج الدراسية المتاحة في كلية الهندسة؟"),
        ("E. 'all' questions", "هات كل المصاريف لكل الكليات بالكامل"),
        ("F. comparison questions", "قارن بين مصاريف طب الأسنان ومصاريف الهندسة"),
        ("G. procedural questions", "ما هي خطوات وسداد المصروفات الدراسية للطلاب الجدد؟")
    ]

    for cat, q in test_queries:
        print(f"\nCategory: {cat}")
        print(f"Query: '{q}'")
        chunks, diag = retriever.retrieve(q, intent="multi_fact" if "all" in cat or "two" in cat else "general_information")
        unique_docs = set([c.filename for c in chunks])
        print(f"Retrieved chunk IDs: {[c.chunk_id for c in chunks[:5]]}")
        print(f"Document names: {list(unique_docs)}")
        print(f"Sections: {list(set([c.section for c in chunks[:5]]))}")
        print(f"Unique relevant sources: {len(unique_docs)}")
        print(f"Latency: {diag['latency_ms']}ms")


def audit_section_5_developers():
    print("\n==================================================")
    print("5. TEST DEVELOPER ENTITY COMPLETENESS")
    print("==================================================")

    devs = DeveloperRegistry.get_all()
    print(f"Total developer entities: {len(devs)}")
    assert len(devs) == 7, "Must be exactly 7 developers!"

    for idx, d in enumerate(devs, 1):
        links = ", ".join([f"{c['type']}: {c['url']}" for c in d.contact_methods]) if d.contact_methods else "None"
        print(f"  {idx}. {d.name} ({d.name_ar}) | Role: {d.role} | Contacts: {links}")

    test_dev_queries = [
        "مين المطورين السبعة؟",
        "مين مسؤول عن backend؟",
        "مين عمل الشات بوت؟",
        "مين مسؤول عن deployment؟",
        "مين مسؤول عن frontend؟",
        "إزاي أتواصل مع Mohammed Mokhtar؟"
    ]

    print("\nDeveloper Queries Evaluation:")
    for q in test_dev_queries:
        if "السبعة" in q or "المطورين" in q and "مين" in q:
            res = DeveloperRegistry.get_all()
        elif "Mokhtar" in q:
            res = DeveloperRegistry.search_by_name("Mokhtar")
        elif "backend" in q:
            res = DeveloperRegistry.search_by_responsibility("backend")
        elif "frontend" in q:
            res = DeveloperRegistry.search_by_responsibility("frontend")
        elif "deployment" in q or "cloud" in q:
            res = DeveloperRegistry.search_by_responsibility("cloud")
        elif "الشات بوت" in q or "chatbot" in q:
            res = DeveloperRegistry.search_by_responsibility("chatbot")
        else:
            res = DeveloperRegistry.get_all()
        
        dev_names = [d.name for d in res]
        print(f"Query: '{q}' -> Matched Devs: {dev_names}")


def audit_section_6_coordination():
    print("\n==================================================")
    print("6. TEST COORDINATION POLICY")
    print("==================================================")

    coord_queries = [
        "الحد الأدنى لكلية الطب كام؟",
        "طب محتاج مجموع كام؟",
        "What is the minimum score for medicine?",
        "هل ده تنسيق جامعة حلوان الأهلية الداخلي؟"
    ]

    for q in coord_queries:
        is_coord = PolicyEngine.is_coordination_query(q)
        formatted = PolicyEngine.apply_coordination_rules("بيان التنسيق الخاص بالطب هو 80%.", q)
        disclaimer_present = COORDINATION_DISCLAIMER in formatted
        print(f"Query: '{q}' | Is Coord Query: {is_coord} | Disclaimer Present: {disclaimer_present}")


def audit_section_7_language():
    print("\n==================================================")
    print("7. TEST LANGUAGE BEHAVIOR")
    print("==================================================")

    lang_queries = [
        "What are the fees for dentistry?",
        "ما هي مصروفات كلية طب الأسنان؟",
        "مصاريف طب الأسنان كام؟",
        "ايه minimum score for medicine?"
    ]

    for q in lang_queries:
        detected = LanguageEngine.detect_language(q)
        print(f"Query: '{q}' -> Lang: {detected['language']} | Style: {detected['style']}")


def audit_section_8_cache():
    print("\n==================================================")
    print("8. CACHE AUDIT")
    print("==================================================")

    cache = CacheManager(cache_file_path=BASE_DIR / "cache" / "test_responses.json")
    
    # Store query
    q1 = "مصاريف كلية الطب"
    q2_en = "What are the fees for medicine?"
    
    cache.put(q1, "إجابة عربية", language="ar")
    cache.put(q2_en, "English response", language="en")

    hit_ar = cache.get(q1, language="ar")
    hit_en = cache.get(q2_en, language="en")

    print(f"Arabic Query Hit: {hit_ar is not None}")
    print(f"English Query Hit: {hit_en is not None}")
    print(f"Knowledge Version: {HNU_KNOWLEDGE_VERSION}")


def build_section_11_evaluation_questions():
    print("\n==================================================")
    print("11. BUILD REALISTIC EVALUATION SET (40 QUESTIONS)")
    print("==================================================")

    questions = [
        # 10 Simple Factual
        {"question": "أين تقع جامعة حلوان الأهلية؟", "category": "simple_factual", "expected_topics": ["location", "address"]},
        {"question": "ما هو البريد الإلكتروني الرسمي للتواصل مع الجامعة؟", "category": "simple_factual", "expected_topics": ["contact", "email"]},
        {"question": "ما هي الأوراق المطلوبة للتقديم في الجامعة؟", "category": "simple_factual", "expected_topics": ["documents", "admission"]},
        {"question": "ما هو رابط موقع التقديم الإلكتروني؟", "category": "simple_factual", "expected_topics": ["link", "admission"]},
        {"question": "هل تمنح الجامعة درجات علمية معتمدة؟", "category": "simple_factual", "expected_topics": ["accreditation"]},
        {"question": "ما هي كليات القطاع الطبي في الجامعة؟", "category": "simple_factual", "expected_topics": ["faculties", "medicine"]},
        {"question": "ما هو نظام الدراسة بجامعة حلوان الأهلية؟", "category": "simple_factual", "expected_topics": ["credit_hours"]},
        {"question": "ما هي الأوراق المطلوبة للطلاب الوافدين؟", "category": "simple_factual", "expected_topics": ["foreign_students"]},
        {"question": "هل تتوفر وسائل مواصلات للجامعة؟", "category": "simple_factual", "expected_topics": ["transportation"]},
        {"question": "ما هي السعة الطلابية لكلية الهندسة؟", "category": "simple_factual", "expected_topics": ["capacity"]},

        # 8 Multi-fact
        {"question": "ما هي مصاريف كلية طب الأسنان وهل ينفع تقسيط؟", "category": "multi_fact", "expected_topics": ["tuition", "installment"]},
        {"question": "ما هي شروط التقديم وما الأوراق المطلوبة للطلاب الجدد؟", "category": "multi_fact", "expected_topics": ["conditions", "documents"]},
        {"question": "ما هي مصاريف الهندسة والحد الأدنى للتنسيق؟", "category": "multi_fact", "expected_topics": ["tuition", "coordination"]},
        {"question": "ما هي برامج كلية علوم الحاسب وما هي مصروفاتها؟", "category": "multi_fact", "expected_topics": ["programs", "tuition"]},
        {"question": "ما هي خطوات وسداد المصروفات وكيفية تأكيد التقديم؟", "category": "multi_fact", "expected_topics": ["payment", "confirmation"]},
        {"question": "ما هي المنح المتاحة وما هي نسبة الخصم للطلاب المتفوقين؟", "category": "multi_fact", "expected_topics": ["scholarships", "discounts"]},
        {"question": "ما هي مصاريف الطب البشري وكيفية التواصل مع شؤون الطلاب؟", "category": "multi_fact", "expected_topics": ["medicine_tuition", "student_affairs"]},
        {"question": "ما هي الكليات المتاحة وما هي مصروفات كل كلية؟", "category": "multi_fact", "expected_topics": ["faculties", "all_tuition"]},

        # 5 Tuition
        {"question": "مصاريف كلية الهندسة كام؟", "category": "tuition", "expected_topics": ["engineering_tuition"]},
        {"question": "مصروفات كلية العلاج الطبيعي كام؟", "category": "tuition", "expected_topics": ["physical_therapy_tuition"]},
        {"question": "كم تكلفة دراسة علوم الحاسب وتكنولوجيا المعلومات؟", "category": "tuition", "expected_topics": ["cs_tuition"]},
        {"question": "ما هي مصاريف كلية الطب والجراحة؟", "category": "tuition", "expected_topics": ["medicine_tuition"]},
        {"question": "ما هي مصاريف كلية العلوم؟", "category": "tuition", "expected_topics": ["science_tuition"]},

        # 5 Admission/Coordination
        {"question": "الحد الأدنى لكلية الطب كام؟", "category": "admission_coordination", "expected_topics": ["medicine_minimum"]},
        {"question": "تنسيق كلية الهندسة كم بالمئة؟", "category": "admission_coordination", "expected_topics": ["engineering_coordination"]},
        {"question": "ما هو الحد الأدنى للقبول بكليات القطاع الطبي؟", "category": "admission_coordination", "expected_topics": ["medical_sector_minimum"]},
        {"question": "ما هو تنسيق كلية علاج طبيعي؟", "category": "admission_coordination", "expected_topics": ["physical_therapy_minimum"]},
        {"question": "هل تنسيق الجامعة الأهلية يختلف عن التنسيق الحكومي؟", "category": "admission_coordination", "expected_topics": ["coordination_disclaimer"]},

        # 5 Developer/Team
        {"question": "مين المطورين السبعة بتوع الموقع الرسميين؟", "category": "developer_team", "expected_topics": ["7_developers"]},
        {"question": "مين مسؤول عن تطوير الـ backend والشات بوت؟", "category": "developer_team", "expected_topics": ["backend_chatbot_dev"]},
        {"question": "مين مسؤول عن تصميم وتطوير الـ frontend؟", "category": "developer_team", "expected_topics": ["frontend_dev"]},
        {"question": "مين مسؤول عن السيرفرات الـ Cloud والـ deployment؟", "category": "developer_team", "expected_topics": ["cloud_deployment_dev"]},
        {"question": "إزاي أتواصل مع المطور Mohammed Mokhtar؟", "category": "developer_team", "expected_topics": ["mokhtar_contact"]},

        # 4 English
        {"question": "What are the fees for dentistry at Helwan National University?", "category": "english", "expected_topics": ["dentistry_tuition"]},
        {"question": "What is the minimum score required for medicine?", "category": "english", "expected_topics": ["medicine_minimum"]},
        {"question": "Where is the university located and how to reach it?", "category": "english", "expected_topics": ["location"]},
        {"question": "Who are the developers of the official website?", "category": "english", "expected_topics": ["developers"]},

        # 3 Egyptian Arabic
        {"question": "مصاريف أسنان كام وهل ينفع ادفع على أقساط؟", "category": "egyptian_arabic", "expected_topics": ["dentistry_tuition", "installment"]},
        {"question": "هو الطب البشري بياخد من مجموع كام؟", "category": "egyptian_arabic", "expected_topics": ["medicine_minimum"]},
        {"question": "ازاي اروح الجامعة بالمواصلات؟", "category": "egyptian_arabic", "expected_topics": ["transportation"]}
    ]

    eval_file = BASE_DIR / "tests" / "evaluation_questions.json"
    eval_file.parent.mkdir(parents=True, exist_ok=True)
    with open(eval_file, "w", encoding="utf-8") as f:
        json.dump(questions, f, ensure_ascii=False, indent=2)

    print(f"✅ Saved {len(questions)} evaluation questions to {eval_file}.")
    return questions


def audit_section_12_retrieval_diagnostics(retriever, questions):
    print("\n==================================================")
    print("12. RETRIEVAL DIAGNOSTICS LOGGING")
    print("==================================================")

    router = QueryRouter()
    context_bld = ContextBuilder()

    diagnostics_log = []
    for item in questions:
        q = item["question"]
        route = router.classify_query(q)
        lang_info = LanguageEngine.detect_language(q)

        start_time = time.time()
        chunks, diag = retriever.retrieve(q, intent=route["intent"], filter_filename=route["filter_filename"])
        latency_ms = diag["latency_ms"]

        context_str = context_bld.build_context(chunks)

        log_entry = {
            "question": q,
            "category": item["category"],
            "intent": route["intent"],
            "language": lang_info["language"],
            "retrieval_latency_ms": latency_ms,
            "retrieved_documents": list(set([c.filename for c in chunks])),
            "retrieved_sections": list(set([c.section for c in chunks])),
            "retrieved_chunk_ids": [c.chunk_id for c in chunks],
            "context_size_chars": len(context_str)
        }
        diagnostics_log.append(log_entry)

    log_file = BASE_DIR / "tests" / "diagnostics_output.json"
    with open(log_file, "w", encoding="utf-8") as f:
        json.dump(diagnostics_log, f, ensure_ascii=False, indent=2)

    print(f"✅ Diagnostics logged for {len(diagnostics_log)} questions to {log_file}.")


def main():
    docs, filenames = audit_section_2_inventory()
    chunks = audit_section_3_chunking(docs)

    index = HybridIndex()
    index.build_index(chunks, enable_sentence_transformers=False)
    retriever = HybridRetriever(index=index)

    audit_section_4_multi_fact(retriever)
    audit_section_5_developers()
    audit_section_6_coordination()
    audit_section_7_language()
    audit_section_8_cache()
    
    questions = build_section_11_evaluation_questions()
    audit_section_12_retrieval_diagnostics(retriever, questions)


if __name__ == "__main__":
    main()