"""
check_sign.py – Kiểm tra và trích xuất thông tin chữ ký số từ PDF.

Tính năng mở rộng:
  - Đọc certificate nhúng trong chữ ký PKCS7/CMS
  - Trích xuất MST từ Subject.user_id ("MST:0311942758")
  - So sánh MST với contract metadata → xác định Bên A / Bên B
"""

from __future__ import annotations
import os
import re
from pyhanko.pdf_utils.reader import PdfFileReader


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _decode_pdf_str(pdf_str) -> str | None:
    """Giải mã PDF string object thành str Python."""
    if not pdf_str:
        return None
    try:
        if hasattr(pdf_str, "decode"):
            return pdf_str.decode("utf-8", "ignore")
        return str(pdf_str)
    except Exception:
        return str(pdf_str)


def _extract_mst_from_subject(subject) -> str:
    """
    Trích xuất MST (Mã số thuế) từ subject của X.509 certificate.

    Certificate VN thường lưu MST trong:
      - user_id  : "MST:0311942758"
      - userid   : "0311942758"
      - serial_number, organization_identifier ...

    Returns MST thuần (10-13 chữ số) hoặc "" nếu không tìm được.
    """
    if subject is None:
        return ""

    # Thử lấy qua dict-like interface (asn1crypto / cryptography)
    potential_keys = [
        "user_id", "userid", "uid",
        "organization_identifier", "serial_number",
        "title", "description",
    ]

    subject_dict: dict = {}
    try:
        subject_dict = dict(subject.native)          # asn1crypto
    except Exception:
        pass

    if not subject_dict:
        try:
            # cryptography library
            for attr in subject:
                oid_name = attr.oid._name if hasattr(attr.oid, "_name") else str(attr.oid)
                subject_dict[oid_name.lower()] = attr.value
        except Exception:
            pass

    for key in potential_keys:
        val = subject_dict.get(key, "")
        if val:
            val_str = str(val).strip()
            # value dạng "MST:0311942758" → bỏ prefix
            if val_str.upper().startswith("MST:"):
                return val_str[4:].strip()
            # value dạng số thuần
            m = re.search(r"\d{10,13}", val_str)
            if m:
                return m.group(0)

    return ""


def _get_signer_cert(sig):
    """
    Trả về certificate của người ký từ embedded_signature của pyhanko.
    Hỗ trợ pyhanko >= 0.20 (sig.signer_cert) và các phiên bản cũ hơn.
    """
    # Thử thuộc tính trực tiếp
    for attr in ("signer_cert", "cert"):
        cert = getattr(sig, attr, None)
        if cert is not None:
            return cert

    # Thử qua signed_data.cert_path
    try:
        sd = sig.signed_data
        certs = list(sd.certs)
        if certs:
            return certs[0]
    except Exception:
        pass

    # Thử parse raw CMS
    try:
        from asn1crypto import cms as asn1_cms, x509 as asn1_x509
        raw = bytes(sig.sig_object["/Contents"])
        content_info = asn1_cms.ContentInfo.load(raw)
        sd = content_info["content"]
        for cert_choice in sd["certificates"]:
            return cert_choice.chosen          # x509.Certificate
    except Exception:
        pass

    return None


# ─────────────────────────────────────────────────────────────────────────────
# Public API
# ─────────────────────────────────────────────────────────────────────────────

def is_pdf_signed(pdf_path: str) -> bool:
    """Kiểm tra xem file PDF đã được ký số hay chưa."""
    if not os.path.exists(pdf_path):
        print(f"Lỗi: Không tìm thấy file '{pdf_path}'")
        return False
    try:
        with open(pdf_path, "rb") as f:
            reader = PdfFileReader(f)
            sigs = reader.embedded_signatures
            return bool(sigs and len(sigs) > 0)
    except Exception as e:
        print(f"Lỗi khi đọc file PDF: {e}")
        return False


def get_signature_info(pdf_path: str,
                       mst_party_a: str = "",
                       mst_party_b: str = "") -> list[dict]:
    """
    Lấy danh sách thông tin chi tiết của các chữ ký bên trong file PDF.
    Trích xuất MST từ Subject certificate và xác định bên ký.

    Args:
        pdf_path:     Đường dẫn file PDF.
        mst_party_a:  MST Bên A (ví dụ "0311942758") để so khớp.
        mst_party_b:  MST Bên B để so khớp.

    Returns:
        list[dict] với các key:
          - field_name  : tên field PDF (BEN_A, BEN_B, ...)
          - name        : tên người ký (CN)
          - company     : tên công ty (O)
          - mst         : mã số thuế trích từ certificate
          - party       : "party_a" | "party_b" | "unknown"
          - date        : ngày ký (chuỗi raw từ PDF)
    """
    result = []

    if not os.path.exists(pdf_path):
        return result

    try:
        with open(pdf_path, "rb") as f:
            reader = PdfFileReader(f)
            signatures = reader.embedded_signatures

            for sig in signatures:
                field_name = sig.field_name or ""
                sig_obj = sig.sig_object.get_object()

                # ── 1. Metadata cơ bản từ PDF dictionary ──────────────────
                name_raw    = _decode_pdf_str(sig_obj.get("/Name"))
                reason_raw  = _decode_pdf_str(sig_obj.get("/Reason"))
                date_raw    = _decode_pdf_str(sig_obj.get("/M"))

                # ── 2. Trích xuất MST từ certificate ───────────────────────
                mst     = ""
                company = ""
                cn      = name_raw or ""

                cert = _get_signer_cert(sig)
                if cert is not None:
                    try:
                        subject = cert.subject
                        mst = _extract_mst_from_subject(subject)

                        # Lấy CN và O từ subject.native nếu có
                        try:
                            native = subject.native
                            cn      = native.get("common_name", cn) or cn
                            company = native.get("organization_name", "") or ""
                        except Exception:
                            pass
                    except Exception as e:
                        print(f"   [check_sign] Lỗi đọc cert subject: {e}")

                # ── 3. Xác định party từ MST ───────────────────────────────
                party = _detect_party(field_name, mst, mst_party_a, mst_party_b)

                result.append({
                    "field_name": field_name,
                    "name":       cn,
                    "company":    company,
                    "mst":        mst,
                    "party":      party,
                    "reason":     reason_raw,
                    "date":       date_raw,
                })

                print(f"   [check_sign] Sig '{field_name}' | MST={mst!r} | party={party}")

    except Exception as e:
        print(f"Lỗi lấy chi tiết chữ ký: {e}")

    return result


def _detect_party(field_name: str,
                  mst: str,
                  mst_party_a: str,
                  mst_party_b: str) -> str:
    """
    Xác định bên ký:
    1/ So sánh MST với mst_party_a / mst_party_b.
    2/ Fallback: đọc field_name (BEN_A, BEN_B, BIEN_A, BIEN_B).
    """
    mst_clean = mst.strip()

    if mst_clean:
        if mst_party_b and mst_clean == mst_party_b.strip():
            return "party_b"
        if mst_party_a and mst_clean == mst_party_a.strip():
            return "party_a"

    # Fallback field_name
    fn_upper = field_name.upper()
    if any(k in fn_upper for k in ("BEN_B", "BIEN_B", "PARTY_B", "BENB")):
        return "party_b"
    if any(k in fn_upper for k in ("BEN_A", "BIEN_A", "PARTY_A", "BENA")):
        return "party_a"

    # Nếu chỉ có 1 trong 2 MST được cung cấp, dùng logic ngược lại
    if mst_clean:
        if mst_party_b and not mst_party_a:
            # Không phải Bên B → Bên A
            return "party_a"
        if mst_party_a and not mst_party_b:
            return "party_b"

    return "unknown"


# ─────────────────────────────────────────────────────────────────────────────
# CLI
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        file_to_check = sys.argv[1]
        mst_a = sys.argv[2] if len(sys.argv) > 2 else ""
        mst_b = sys.argv[3] if len(sys.argv) > 3 else ""
    else:
        file_to_check = input("Nhập đường dẫn file PDF: ").strip(' "\'')
        mst_a = input("MST Bên A (Enter bỏ qua): ").strip()
        mst_b = input("MST Bên B (Enter bỏ qua): ").strip()

    if not file_to_check:
        sys.exit(0)

    signed = is_pdf_signed(file_to_check)
    if signed:
        print(f"\n✅ File '{os.path.basename(file_to_check)}' ĐÃ ĐƯỢC KÝ SỐ.")
        details = get_signature_info(file_to_check, mst_a, mst_b)
        if details:
            print(f"\n👉 Tìm thấy {len(details)} chữ ký:")
            for i, d in enumerate(details, 1):
                party_label = {"party_a": "Bên A", "party_b": "Bên B"}.get(d["party"], "Không rõ")
                print(f"\n   [{i}] Field  : {d['field_name']}")
                print(f"       Người ký: {d['name']}")
                if d["company"]: print(f"       Công ty : {d['company']}")
                print(f"       MST    : {d['mst'] or 'N/A'}")
                print(f"       Bên ký : {party_label}")
                if d["date"]:   print(f"       Ngày ký : {d['date']}")
    else:
        print(f"\n❌ File '{os.path.basename(file_to_check)}' CHƯA CÓ chữ ký số.")
