"""
batch_test.py - Test tất cả file DOCX/DOC trong thư mục test/
Chạy: python test/batch_test.py
"""
import sys
import io
import os
import re
import time

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
ROOT = os.path.dirname(os.path.dirname(__file__))
sys.path.insert(0, ROOT)

from docx_converter import convert_docx_to_markdown
from ocr_drive import get_drive_service, ocr_file
from text_to_markdown import text_to_markdown

TEST_DIR = os.path.join(ROOT, "test")
RESULTS_DIR = os.path.join(ROOT, "test_results")
os.makedirs(RESULTS_DIR, exist_ok=True)

MIME_MAP = {
    ".doc": "application/msword",
    ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    ".pdf": "application/pdf",
}


def count_dieu(md):
    return len(re.findall(r"(?:ĐIỀU|Điều)\s+\d+", md, re.UNICODE))


def count_headings(md):
    return sum(1 for l in md.split("\n") if l.startswith("#"))


def get_top_headings(md, n=8):
    return [l for l in md.split("\n") if l.startswith("#")][:n]


def test_docx(fpath, fname, drive):
    """Test DOCX: chạy 3 strategy, so sánh."""
    with open(fpath, "rb") as f:
        file_bytes = f.read()

    results = {}

    # Strategy A: docx_converter native
    try:
        res = convert_docx_to_markdown(file_bytes, fname)
        md_a = res.get("document", {}).get("md_content", "")
        results["docx_native"] = md_a
    except Exception as e:
        print(f"    [A] docx_native FAIL: {e}")
        results["docx_native"] = ""

    # Strategy B: Drive → export DOCX → docx_converter
    try:
        docx_bytes = ocr_file(drive, fpath, source_mime_type=MIME_MAP[".docx"],
                              filename=fname, export_mode="docx")
        if docx_bytes:
            res = convert_docx_to_markdown(docx_bytes, fname)
            md_b = res.get("document", {}).get("md_content", "")
            results["drive_docx"] = md_b
        else:
            results["drive_docx"] = ""
    except Exception as e:
        print(f"    [B] drive_docx FAIL: {e}")
        results["drive_docx"] = ""

    # Strategy C: Drive → export text → text_to_markdown
    try:
        raw = ocr_file(drive, fpath, source_mime_type=MIME_MAP[".docx"],
                       filename=fname, export_mode="text")
        md_c = text_to_markdown(raw) if raw else ""
        results["drive_text"] = md_c
    except Exception as e:
        print(f"    [C] drive_text FAIL: {e}")
        results["drive_text"] = ""

    return results


def test_doc(fpath, fname, drive):
    """Test DOC: chạy 2 strategy (python-docx không đọc .doc)."""
    results = {}

    # Strategy B: Drive → export DOCX → docx_converter
    try:
        docx_bytes = ocr_file(drive, fpath, source_mime_type=MIME_MAP[".doc"],
                              filename=fname, export_mode="docx")
        if docx_bytes:
            res = convert_docx_to_markdown(docx_bytes, fname)
            md_b = res.get("document", {}).get("md_content", "")
            results["drive_docx"] = md_b
        else:
            results["drive_docx"] = ""
    except Exception as e:
        print(f"    [B] drive_docx FAIL: {e}")
        results["drive_docx"] = ""

    # Strategy C: Drive → export text → text_to_markdown
    try:
        raw = ocr_file(drive, fpath, source_mime_type=MIME_MAP[".doc"],
                       filename=fname, export_mode="text")
        md_c = text_to_markdown(raw) if raw else ""
        results["drive_text"] = md_c
    except Exception as e:
        print(f"    [C] drive_text FAIL: {e}")
        results["drive_text"] = ""

    return results


def pick_best(results):
    """Chọn strategy nhiều 'Điều' nhất."""
    best_key = ""
    best_md = ""
    best_dieu = -1
    for key, md in results.items():
        d = count_dieu(md)
        if d > best_dieu or (d == best_dieu and len(md) > len(best_md)):
            best_key = key
            best_md = md
            best_dieu = d
    return best_key, best_md


# ─────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────

print("=" * 70)
print("  BATCH TEST — DOCX/DOC → Markdown (Focus: 'Điều' extraction)")
print("=" * 70)

drive = get_drive_service()
files = sorted(os.listdir(TEST_DIR))
summary = []

for fname in files:
    fpath = os.path.join(TEST_DIR, fname)
    if not os.path.isfile(fpath):
        continue

    ext = os.path.splitext(fname)[1].lower()
    if ext not in (".doc", ".docx"):
        continue

    print(f"\n{'─'*70}")
    print(f"  📄 {fname} ({os.path.getsize(fpath)//1024} KB)")

    t0 = time.time()
    try:
        if ext == ".docx":
            results = test_docx(fpath, fname, drive)
        else:
            results = test_doc(fpath, fname, drive)

        elapsed = time.time() - t0

        # Print mỗi strategy
        for key, md in results.items():
            d = count_dieu(md)
            h = count_headings(md)
            flag = "✅" if d > 0 else "⚠️"
            print(f"    {flag} {key:15s} → {len(md):>6,} chars | {h:>3} headings | {d:>2} 'Điều'")

        # Pick best
        best_key, best_md = pick_best(results)
        best_dieu = count_dieu(best_md)
        print(f"    🏆 BEST: {best_key} ({best_dieu} Điều, {len(best_md):,} chars) [{elapsed:.1f}s]")

        # Save best result
        out_name = os.path.splitext(fname)[0] + "_result.md"
        out_path = os.path.join(RESULTS_DIR, out_name)
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(f"<!-- Backend: {best_key} | Điều: {best_dieu} -->\n")
            f.write(f"# {fname}\n\n")
            f.write(best_md)

        # Top headings
        top = get_top_headings(best_md, 5)
        if top:
            print(f"    Top headings:")
            for h in top:
                print(f"      {h[:80]}")

        summary.append((fname, best_key, best_dieu, len(best_md)))

    except Exception as e:
        elapsed = time.time() - t0
        print(f"    ❌ ERROR ({elapsed:.1f}s): {e}")
        summary.append((fname, "ERROR", 0, 0))

# ── Summary ──
print(f"\n{'='*70}")
print(f"  SUMMARY ({len(summary)} files)")
print(f"{'='*70}")
ok = [s for s in summary if s[1] != "ERROR"]
no_dieu = [s for s in ok if s[2] == 0]
has_dieu = [s for s in ok if s[2] > 0]
print(f"  ✅ Có 'Điều':     {len(has_dieu)} file")
print(f"  ⚠️  Không 'Điều': {len(no_dieu)} file")
print(f"  ❌ Lỗi:          {len(summary) - len(ok)} file")
for s in summary:
    flag = "✅" if s[2] > 0 else ("❌" if s[1] == "ERROR" else "⚠️")
    print(f"  {flag} {s[0]:50s} → {s[1]:15s} | {s[2]:>2} Điều | {s[3]:>6,} chars")
print(f"\n  Results: {RESULTS_DIR}/")
print("=" * 70)
