"""Chấm bài tự động (giống Themis, bản đơn giản) cho lớp học.

Cách dùng:
    python cham_bai.py <bai_lam.py> <thu_muc_test> [-t 1.0] [-d 10] [--file TENBAI]

Thư mục test có thể theo 1 trong 2 kiểu:
    tests/01.inp, tests/01.out, tests/02.inp, ...         (kiểu đơn giản)
    tests/Test01/TENBAI.INP, tests/Test01/TENBAI.OUT, ... (kiểu Themis)

Mặc định bài làm đọc bàn phím (stdin) và in ra màn hình (stdout).
Nếu đề yêu cầu đọc/ghi file TENBAI.INP / TENBAI.OUT thì thêm --file TENBAI.
So sánh kết quả: bỏ qua khác biệt về dấu cách và xuống dòng.
"""

import argparse
import os
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path


def tim_test(thu_muc):
    """Trả về danh sách (tên test, file inp, file out)."""
    ds = []
    for p in sorted(thu_muc.iterdir()):
        if p.is_file() and p.suffix.lower() == ".inp":
            out = next((q for q in p.parent.glob(p.stem + ".*") if q.suffix.lower() == ".out"), None)
            if out:
                ds.append((p.stem, p, out))
        elif p.is_dir():
            inp = next((q for q in p.iterdir() if q.suffix.lower() == ".inp"), None)
            out = next((q for q in p.iterdir() if q.suffix.lower() == ".out"), None)
            if inp and out:
                ds.append((p.name, inp, out))
    return ds


def doc(path):
    return Path(path).read_text(encoding="utf-8", errors="replace")


def rut_gon(s, n=60):
    s = " ".join(s.split())
    return s if len(s) <= n else s[:n] + "..."


def chay_mot_test(bai_lam, inp, gioi_han, ten_file):
    """Chạy bài làm trên một test. Trả về (trạng thái, thời gian, output, thông báo lỗi)."""
    env = dict(os.environ, PYTHONIOENCODING="utf-8")
    with tempfile.TemporaryDirectory() as tmp:
        if ten_file:
            shutil.copy(inp, Path(tmp) / f"{ten_file}.INP")
            du_lieu_vao = ""
        else:
            du_lieu_vao = doc(inp)
        bat_dau = time.perf_counter()
        try:
            kq = subprocess.run(
                [sys.executable, str(bai_lam)],
                input=du_lieu_vao, capture_output=True, text=True,
                encoding="utf-8", errors="replace",
                timeout=gioi_han, cwd=tmp, env=env,
            )
        except subprocess.TimeoutExpired:
            return "TLE", gioi_han, "", ""
        tg = time.perf_counter() - bat_dau
        if kq.returncode != 0:
            loi = kq.stderr.strip().splitlines()
            return "RE", tg, "", loi[-1] if loi else f"mã thoát {kq.returncode}"
        if ten_file:
            f = next((q for q in Path(tmp).iterdir() if q.name.upper() == f"{ten_file.upper()}.OUT"), None)
            output = doc(f) if f else ""
        else:
            output = kq.stdout
        return "OK", tg, output, ""


def main():
    ap = argparse.ArgumentParser(description="Chấm bài tự động bằng bộ test.")
    ap.add_argument("bai_lam", help="file bài làm .py")
    ap.add_argument("tests", help="thư mục chứa test")
    ap.add_argument("-t", "--time", type=float, default=1.0, help="giới hạn thời gian mỗi test (giây), mặc định 1")
    ap.add_argument("-d", "--diem", type=float, default=10, help="tổng điểm của bài, mặc định 10")
    ap.add_argument("--file", metavar="TENBAI", help="đọc TENBAI.INP, ghi TENBAI.OUT thay cho bàn phím/màn hình")
    args = ap.parse_args()

    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")

    bai_lam = Path(args.bai_lam).resolve()
    thu_muc = Path(args.tests)
    if not bai_lam.is_file():
        sys.exit(f"Không tìm thấy file bài làm: {bai_lam}")
    if not thu_muc.is_dir():
        sys.exit(f"Không tìm thấy thư mục test: {thu_muc}")
    ds_test = tim_test(thu_muc)
    if not ds_test:
        sys.exit(f"Không có test nào trong {thu_muc}")

    nhan = {"AC": "ĐÚNG", "WA": "SAI KẾT QUẢ", "TLE": "QUÁ THỜI GIAN", "RE": "LỖI KHI CHẠY"}
    so_dung = 0
    print(f"Chấm {bai_lam.name} – {len(ds_test)} test – giới hạn {args.time:g}s/test\n")
    for ten, inp, out in ds_test:
        trang_thai, tg, output, loi = chay_mot_test(bai_lam, inp, args.time, args.file)
        mong_doi = doc(out)
        if trang_thai == "OK":
            trang_thai = "AC" if output.split() == mong_doi.split() else "WA"
        so_dung += trang_thai == "AC"
        print(f"  {ten:<10} {nhan[trang_thai]:<15} {tg:.2f}s")
        if trang_thai == "WA":
            print(f"             cần in : {rut_gon(mong_doi)}")
            print(f"             em in  : {rut_gon(output) or '(không in gì)'}")
        elif trang_thai == "RE":
            print(f"             {loi}")

    diem = args.diem * so_dung / len(ds_test)
    print(f"\nKết quả: đúng {so_dung}/{len(ds_test)} test  →  {diem:g}/{args.diem:g} điểm")


if __name__ == "__main__":
    main()