#!/usr/bin/env python3
"""
Cross-language RAG evaluation loop.

This runner reuses the normal E2E RAG eval flow, but adds a small target-language
check for cases where the user asks in one language and requests the answer in
another language.
"""

from __future__ import annotations

import argparse
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List

from rag_eval_loop import (
    DEFAULT_ABSTAIN_PHRASE,
    DEFAULT_POLL_INTERVAL_SECONDS,
    DEFAULT_TIMEOUT_SECONDS,
    REPO_ROOT,
    EvalCaseResult,
    EvalRunner,
    _content_tokens,
    _number_claims,
    _sentences,
    _source_text,
    load_eval_file,
    normalize,
    resolve_doc_paths,
    score_answer_correctness,
    score_retrieval_recall,
)


DEFAULT_EVAL_FILE = Path(__file__).with_name("rag_eval_cross_language.json")

_GENERAL_LANGUAGE_MARKERS = {
    "en": {
        "the", "is", "are", "for", "with", "and", "scheduled", "required",
        "maximum", "minimum", "course", "lecturer", "office", "hour",
        "assignment", "submission", "deadline", "format", "weight", "worth",
    },
    "id": {
        "adalah", "untuk", "dalam", "yang", "pukul", "ruang", "tanggal",
        "nilai", "batas", "hari", "jadwal", "dengan", "melalui", "ukuran",
        "harus", "dosen", "sesi", "berkontribusi", "terhadap", "memiliki",
        "bobot", "ditetapkan", "diperbolehkan", "diizinkan", "digunakan",
        "siswa", "mahasiswa", "buku", "teks", "inti", "sebesar", "denda",
        "keterlambatan", "ambang", "pengumpulan", "kelompok", "nama",
    },
}

_BILINGUAL_TOKEN_GROUPS = [
    {"course", "mata", "kuliah"},
    {"lecturer", "instructor", "dosen", "pengajar"},
    {"teaching", "assistant", "asisten"},
    {"office", "hour", "jam"},
    {"room", "ruang"},
    {"held", "diadakan", "berlangsung"},
    {"schedule", "jadwal"},
    {"monday", "senin"},
    {"tuesday", "selasa"},
    {"wednesday", "rabu"},
    {"thursday", "kamis"},
    {"friday", "jumat"},
    {"quiz", "quizzes", "kuis"},
    {"minutes", "menit"},
    {"lab", "praktikum"},
    {"assessment", "grading", "grade", "penilaian", "nilai"},
    {"assignment", "assignments", "tugas"},
    {"midterm", "uts", "ujian", "tengah", "semester"},
    {"final", "project", "demo", "proyek", "akhir"},
    {"minimum", "minimal"},
    {"passing", "pass", "lulus"},
    {"score", "nilai"},
    {"deadline", "batas", "tenggat", "waktu"},
    {"submission", "submitted", "pengumpulan", "dikumpulkan"},
    {"via", "through", "melalui", "lewat"},
    {"exam", "ujian"},
    {"date", "tanggal"},
    {"public", "holiday", "libur"},
    {"class", "kelas"},
    {"make", "up", "make-up", "pengganti"},
    {"format", "format"},
    {"pdf"},
    {"maximum", "maksimum"},
    {"file", "berkas"},
    {"size", "ukuran"},
    {"required", "wajib", "diwajibkan", "harus"},
    {"filename", "nama", "file"},
    {"pattern", "pola"},
    {"late", "terlambat", "keterlambatan"},
    {"penalty", "penalti", "denda"},
    {"window", "batas", "periode"},
    {"days", "hari"},
    {"plagiarism", "plagiarisme"},
    {"threshold", "ambang", "batas"},
    {"similarity", "kesamaan"},
    {"below", "bawah"},
    {"group", "kelompok"},
    {"students", "mahasiswa", "siswa"},
    {"branch"},
    {"naming", "penamaan"},
    {"commit"},
    {"message", "pesan"},
    {"convention", "konvensi"},
    {"programming", "pemrograman"},
    {"language", "bahasa"},
    {"framework", "kerangka"},
    {"core", "inti", "utama"},
    {"textbook", "book", "buku", "teks"},
    {"edition", "edisi"},
]


def _score_target_language(answer: str, case: Dict[str, Any]) -> List[str]:
    expected_language = (case.get("expected_language") or "").strip().lower()
    if not expected_language:
        return []

    markers = [
        normalize(marker)
        for marker in case.get("language_expected_any", [])
        if marker
    ]
    if not markers:
        return []

    answer_norm = normalize(answer)
    if any(marker in answer_norm for marker in markers):
        return []
    if _looks_like_language(answer_norm, expected_language):
        return []

    label = {
        "en": "English",
        "id": "Bahasa Indonesia",
    }.get(expected_language, expected_language)
    return [f"language: expected {label} marker, found none of {markers}"]


def _looks_like_language(answer_norm: str, expected_language: str) -> bool:
    markers = _GENERAL_LANGUAGE_MARKERS.get(expected_language, set())
    if not markers:
        return False

    tokens = set(answer_norm.split())
    return len(tokens & markers) >= 2


def _expand_with_bilingual_equivalents(tokens: set[str]) -> set[str]:
    expanded = set(tokens)
    for group in _BILINGUAL_TOKEN_GROUPS:
        if tokens & group:
            expanded.update(group)
    return expanded


def _score_cross_language_faithfulness(
    answer: str,
    sources: List[Dict[str, Any]],
    abstain_phrase: str,
    min_overlap: float,
) -> tuple[bool, List[str]]:
    answer_norm = normalize(answer)
    if not answer_norm or abstain_phrase in answer_norm:
        return True, []

    source_norm = normalize(_source_text(sources))
    if not source_norm:
        return False, ["faithfulness cannot be checked because no sources were returned"]

    source_tokens = _expand_with_bilingual_equivalents(set(_content_tokens(source_norm)))
    source_numbers = _number_claims(source_norm)
    reasons: List[str] = []

    for sentence in _sentences(answer):
        sentence_tokens = _content_tokens(sentence)
        sentence_numbers = _number_claims(sentence)
        if len(sentence_tokens) < 4 and not sentence_numbers:
            continue

        missing_numbers = sorted(sentence_numbers - source_numbers)
        if missing_numbers:
            reasons.append(
                f"faithfulness: sentence has numbers not found in sources: {', '.join(missing_numbers)}"
            )
            continue

        overlap = sum(1 for token in sentence_tokens if token in source_tokens)
        ratio = overlap / max(len(sentence_tokens), 1)
        if ratio < min_overlap:
            reasons.append(
                f"faithfulness: low bilingual source overlap ({ratio:.2f}) for sentence: {sentence}"
            )

    return len(reasons) == 0, reasons


def _score_cross_language_case(
    answer: str,
    sources: List[Dict[str, Any]],
    case: Dict[str, Any],
    abstain_phrase: str,
    idx: int,
    faithfulness_min_overlap: float,
) -> EvalCaseResult:
    answer_reasons = score_answer_correctness(answer, case, abstain_phrase)
    retrieval_checked, retrieval_passed, retrieval_reasons = score_retrieval_recall(sources, case)
    faithfulness_passed, faithfulness_reasons = _score_cross_language_faithfulness(
        answer,
        sources,
        abstain_phrase,
        faithfulness_min_overlap,
    )

    reasons = [
        *(f"answer: {reason}" for reason in answer_reasons),
        *(f"retrieval: {reason}" for reason in retrieval_reasons),
        *(f"faithfulness: {reason}" for reason in faithfulness_reasons),
    ]
    answer_passed = len(answer_reasons) == 0

    return EvalCaseResult(
        index=idx,
        question=case.get("question", ""),
        passed=answer_passed and retrieval_passed and faithfulness_passed,
        answer_passed=answer_passed,
        retrieval_checked=retrieval_checked,
        retrieval_passed=retrieval_passed,
        faithfulness_passed=faithfulness_passed,
        reasons=reasons,
        answer_text=(answer or ""),
        sources=sources,
    )


def _print_case_result(result: EvalCaseResult, max_preview_chars: int, show_full_answers: bool) -> None:
    verdict = "PASS" if result.passed else "FAIL"
    print(f"[{verdict}] Q{result.index}: {result.question}")
    print(
        "  metrics: "
        f"answer={'PASS' if result.answer_passed else 'FAIL'}, "
        f"retrieval={'PASS' if result.retrieval_passed else 'FAIL'}"
        f"{'' if result.retrieval_checked else ' (not checked)'}, "
        f"faithfulness={'PASS' if result.faithfulness_passed else 'FAIL'}, "
        f"sources={len(result.sources)}"
    )
    if not result.passed:
        for reason in result.reasons:
            print(f"  - {reason}")

    shown_answer = result.answer_text if show_full_answers else result.answer_text[:max_preview_chars]
    print(f"  answer: {shown_answer}")


def main() -> int:
    parser = argparse.ArgumentParser(description="Run cross-language RAG E2E evaluation")
    parser.add_argument("--eval-file", default=str(DEFAULT_EVAL_FILE), help="Path to cross-language eval JSON file")
    parser.add_argument("--base-url", default="http://127.0.0.1:8000", help="Server base URL")
    parser.add_argument("--api-prefix", default="/api/v1", help="API prefix")
    parser.add_argument("--password", default="Pass1234!", help="Password for generated test users")
    parser.add_argument("--request-timeout", type=int, default=60, help="Per-request timeout seconds")
    parser.add_argument("--question-timeout", type=int, default=240, help="Timeout seconds for each chat answer request")
    parser.add_argument("--ingest-timeout", type=int, default=DEFAULT_TIMEOUT_SECONDS, help="Ingestion wait timeout seconds")
    parser.add_argument("--poll-interval", type=int, default=DEFAULT_POLL_INTERVAL_SECONDS, help="Polling interval seconds")
    parser.add_argument("--abstain-phrase", default=DEFAULT_ABSTAIN_PHRASE, help="Expected abstention phrase")
    parser.add_argument(
        "--faithfulness-min-overlap",
        type=float,
        default=0.35,
        help="Minimum content-token overlap between each answer sentence and retrieved sources.",
    )
    parser.add_argument("--max-preview-chars", type=int, default=240, help="Max characters shown per answer")
    parser.add_argument("--show-full-answers", action="store_true", help="Print full answer text for every case")
    parser.add_argument(
        "--reuse-course-id",
        action="store_true",
        help="Use the exact course_id from the eval file instead of appending a timestamp.",
    )
    parser.add_argument(
        "--email-domain",
        default="sevima.co.id",
        help="Domain for temporary eval users. The backend rejects placeholder domains such as example.com.",
    )
    args = parser.parse_args()

    payload = load_eval_file(Path(args.eval_file).resolve())
    ts = datetime.utcnow().strftime("%Y%m%d%H%M%S")
    email_domain = args.email_domain.strip().lower().lstrip("@")

    runner = EvalRunner(
        base_url=args.base_url,
        api_prefix=args.api_prefix,
        request_timeout=args.request_timeout,
    )

    print("[INFO] Starting cross-language RAG E2E evaluation loop")
    print(f"[INFO] Base URL: {args.base_url}{args.api_prefix}")

    lecturer_token = runner.ensure_user(
        role="lecturer",
        email=f"rag_cross_eval_lecturer_{ts}@{email_domain}",
        password=args.password,
        name="RAG Cross-Language Eval Lecturer",
        identity_number=f"198502{ts[-12:]}",
    )
    student_token = runner.ensure_user(
        role="student",
        email=f"rag_cross_eval_student_{ts}@{email_domain}",
        password=args.password,
        name="RAG Cross-Language Eval Student",
        identity_number=f"{ts[-10:]}",
    )

    eval_course_id = payload["course_id"]
    course_id = eval_course_id if args.reuse_course_id else f"{eval_course_id}-{ts}"
    print(f"[INFO] Eval course_id: {eval_course_id}")
    print(f"[INFO] Run course_id: {course_id}")

    docs = resolve_doc_paths(REPO_ROOT, payload["documents"])
    uploaded_doc_ids = runner.upload_documents(
        token=lecturer_token,
        course_id=course_id,
        document_paths=docs,
    )

    statuses = runner.wait_for_ingestion(
        token=student_token,
        course_id=course_id,
        uploaded_doc_ids=uploaded_doc_ids,
        timeout_seconds=args.ingest_timeout,
        poll_interval_seconds=args.poll_interval,
    )

    failed = [doc_id for doc_id, details in statuses.items() if details.get("status") == "failed"]
    if failed:
        print(f"[ERROR] Ingestion failed for document_ids: {failed}")
        for doc_id in failed:
            details = statuses.get(doc_id, {})
            print(f"  - doc_id={doc_id} filename={details.get('filename') or '<unknown>'}")
            print(f"    error={details.get('error') or '<no error message provided by backend>'}")
        return 2

    print("[INFO] All uploaded documents are ready")

    session_id = runner.create_student_session(
        token=student_token,
        course_id=course_id,
        title=f"RAG Cross-Language Eval Session {ts}",
    )

    results: List[EvalCaseResult] = []
    for idx, case in enumerate(payload["questions"], start=1):
        answer_payload = runner.ask_question(
            token=student_token,
            session_id=session_id,
            question=case.get("question", ""),
            question_timeout_seconds=args.question_timeout,
        )
        language_reasons = _score_target_language(answer_payload["content"], case)
        result = _score_cross_language_case(
            answer_payload["content"],
            answer_payload["sources"],
            case,
            normalize(args.abstain_phrase),
            idx,
            args.faithfulness_min_overlap,
        )
        if language_reasons:
            result.reasons.extend(language_reasons)
            result.answer_passed = False
            result.passed = False

        results.append(result)
        _print_case_result(result, args.max_preview_chars, args.show_full_answers)

    total = len(results)
    passed = sum(1 for result in results if result.passed)
    failed_count = total - passed
    answer_passed = sum(1 for result in results if result.answer_passed)
    retrieval_checked = sum(1 for result in results if result.retrieval_checked)
    retrieval_passed = sum(1 for result in results if result.retrieval_checked and result.retrieval_passed)
    faithfulness_passed = sum(1 for result in results if result.faithfulness_passed)

    print("\n=== CROSS-LANGUAGE RAG E2E SUMMARY ===")
    print(f"total_cases: {total}")
    print(f"overall_passed: {passed}")
    print(f"overall_failed: {failed_count}")
    print(f"overall_score_percent: {(passed / total) * 100 if total else 0.0:.1f}")
    print(f"answer_and_language_passed: {answer_passed}")
    print(f"retrieval_recall_checked_cases: {retrieval_checked}")
    print(f"retrieval_recall_passed: {retrieval_passed}")
    print(f"faithfulness_passed: {faithfulness_passed}")

    return 0 if failed_count == 0 else 1


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        print("\n[INFO] Interrupted by user")
        raise SystemExit(130)