import ctypes
import hashlib as _hashlib
from ctypes import wintypes
import datetime
import hashlib
from pyhanko.sign.signers import Signer
from pyhanko.pdf_utils import misc
from asn1crypto import x509 as asn1_x509

# Windows API Constants
CERT_KEY_PROV_INFO_PROP_ID = 2
PKCS_7_ASN_ENCODING = 0x00010000
X509_ASN_ENCODING = 0x00000001
MY_ENCODING_TYPE = PKCS_7_ASN_ENCODING | X509_ASN_ENCODING

crypt32 = ctypes.WinDLL('crypt32.dll')
cryptui = ctypes.WinDLL('cryptui.dll')
ncrypt = ctypes.WinDLL('ncrypt.dll')
advapi32 = ctypes.WinDLL('advapi32.dll')

# Hằng số và Cấu trúc bổ sung
PROV_RSA_AES = 24
CRYPT_KEY_PROV_INFO_PROP_ID = 2
CERT_STORE_PROV_MEMORY = 2
CERT_STORE_CREATE_NEW_FLAG = 0x00002000
CERT_STORE_ADD_ALWAYS = 4

class CERT_CONTEXT(ctypes.Structure):
    _fields_ = [
        ("dwCertEncodingType", wintypes.DWORD),
        ("pbCertEncoded", ctypes.POINTER(ctypes.c_byte)),
        ("cbCertEncoded", wintypes.DWORD),
        ("pCertInfo", ctypes.c_void_p),
        ("hCertStore", ctypes.c_void_p),
    ]

class CRYPT_KEY_PROV_INFO(ctypes.Structure):
    _fields_ = [
        ("pwszContainerName", wintypes.LPCWSTR),
        ("pwszProvName", wintypes.LPCWSTR),
        ("dwProvType", wintypes.DWORD),
        ("dwFlags", wintypes.DWORD),
        ("cProvParam", wintypes.DWORD),
        ("rgProvParam", ctypes.c_void_p),
        ("dwKeySpec", wintypes.DWORD),
    ]

crypt32.CertGetCertificateContextProperty.argtypes = [
    ctypes.c_void_p, wintypes.DWORD, ctypes.c_void_p, ctypes.POINTER(wintypes.DWORD)
]

# Định nghĩa kiểu dữ liệu cho các hàm để tránh lỗi Access Violation và Overflow
crypt32.CertOpenSystemStoreW.argtypes = [ctypes.c_void_p, wintypes.LPCWSTR]
crypt32.CertOpenSystemStoreW.restype = ctypes.c_void_p # HCERTSTORE

crypt32.CertOpenStore.argtypes = [
    ctypes.c_void_p, wintypes.DWORD, ctypes.c_void_p, wintypes.DWORD, ctypes.c_void_p
]
crypt32.CertOpenStore.restype = ctypes.c_void_p

crypt32.CertAddCertificateContextToStore.argtypes = [
    ctypes.c_void_p, ctypes.c_void_p, wintypes.DWORD, ctypes.c_void_p
]
crypt32.CertAddCertificateContextToStore.restype = wintypes.BOOL

cryptui.CryptUIDlgSelectCertificateFromStore.argtypes = [
    ctypes.c_void_p, # hCertStore
    ctypes.c_void_p, # hwnd
    wintypes.LPCWSTR, # pwszTitle
    wintypes.LPCWSTR, # pwszDisplayString
    wintypes.DWORD,   # dwFlags
    wintypes.DWORD,   # dwDontUseUsed
    ctypes.c_void_p   # pvReserved
]
cryptui.CryptUIDlgSelectCertificateFromStore.restype = ctypes.c_void_p

crypt32.CertFreeCertificateContext.argtypes = [ctypes.c_void_p]
crypt32.CertFreeCertificateContext.restype = wintypes.BOOL

crypt32.CertDuplicateCertificateContext.argtypes = [ctypes.c_void_p]
crypt32.CertDuplicateCertificateContext.restype = ctypes.c_void_p

crypt32.CertCloseStore.argtypes = [ctypes.c_void_p, wintypes.DWORD]
crypt32.CertCloseStore.restype = wintypes.BOOL

# Dùng để duyệt qua tất cả cert trong một store (lấy CA chain)
crypt32.CertEnumCertificatesInStore.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
crypt32.CertEnumCertificatesInStore.restype = ctypes.c_void_p

crypt32.CryptAcquireCertificatePrivateKey.argtypes = [
    ctypes.c_void_p, # pCertContext
    wintypes.DWORD,  # dwFlags
    ctypes.c_void_p, # pvReserved
    ctypes.POINTER(ctypes.c_void_p), # phCryptProvOrNCryptKey
    ctypes.POINTER(wintypes.DWORD),  # pdwKeySpec
    ctypes.POINTER(wintypes.BOOL)    # pfCallerFreeProvOrNCryptKey
]
crypt32.CryptAcquireCertificatePrivateKey.restype = wintypes.BOOL

ncrypt.NCryptSignHash.argtypes = [
    ctypes.c_void_p, # hKey
    ctypes.c_void_p, # pPaddingInfo
    ctypes.c_char_p, # pbHashValue
    wintypes.DWORD,  # cbHashValue
    ctypes.c_char_p, # pbSignature
    wintypes.DWORD,  # cbSignature
    ctypes.POINTER(wintypes.DWORD), # pcbResult
    wintypes.DWORD   # dwFlags
]
ncrypt.NCryptSignHash.restype = wintypes.DWORD

ncrypt.NCryptFreeObject.argtypes = [ctypes.c_void_p]
ncrypt.NCryptFreeObject.restype = wintypes.DWORD

# NCrypt Flags
NCRYPT_PAD_PKCS1_FLAG = 0x00000002

class BCRYPT_PKCS1_PADDING_INFO(ctypes.Structure):
    _fields_ = [("pszAlgId", wintypes.LPCWSTR)]


def _cert_from_context(cert_context_ptr):
    ctx = ctypes.cast(cert_context_ptr, ctypes.POINTER(CERT_CONTEXT)).contents
    if not ctx.pbCertEncoded or ctx.cbCertEncoded <= 0:
        return None
    der = ctypes.string_at(ctx.pbCertEncoded, ctx.cbCertEncoded)
    return asn1_x509.Certificate.load(der)


def _normalize_dt(value):
    if value.tzinfo is None:
        return value.replace(tzinfo=datetime.timezone.utc)
    return value.astimezone(datetime.timezone.utc)


def _is_asn1_cert_currently_valid(cert):
    validity = cert["tbs_certificate"]["validity"]
    not_before = _normalize_dt(validity["not_before"].native)
    not_after = _normalize_dt(validity["not_after"].native)
    now = datetime.datetime.now(datetime.timezone.utc)
    if now < not_before:
        return False, f"chưa có hiệu lực đến {not_before.strftime('%d/%m/%Y %H:%M')}"
    if now > not_after:
        return False, f"hết hạn từ {not_after.strftime('%d/%m/%Y %H:%M')}"
    return True, f"còn hạn đến {not_after.strftime('%d/%m/%Y %H:%M')}"


def _cert_context_has_private_key(cert_context_ptr):
    size = wintypes.DWORD()
    return bool(
        crypt32.CertGetCertificateContextProperty(
            cert_context_ptr,
            CERT_KEY_PROV_INFO_PROP_ID,
            None,
            ctypes.byref(size),
        )
        and size.value > 0
    )


def _build_valid_certificate_store(source_store):
    valid_store = crypt32.CertOpenStore(
        ctypes.c_void_p(CERT_STORE_PROV_MEMORY),
        0,
        None,
        CERT_STORE_CREATE_NEW_FLAG,
        None,
    )
    if not valid_store:
        return None, 0

    valid_count = 0
    p_ctx = None
    while True:
        p_ctx = crypt32.CertEnumCertificatesInStore(source_store, p_ctx)
        if not p_ctx:
            break
        try:
            if not _cert_context_has_private_key(p_ctx):
                continue
            cert = _cert_from_context(p_ctx)
            if cert is None:
                continue
            is_valid, reason = _is_asn1_cert_currently_valid(cert)
            subject = cert.subject.native
            cn = subject.get("common_name", "Unknown")
            if not is_valid:
                print(f"   [CertFilter] Bỏ qua chứng thư '{cn}': {reason}")
                continue
            if crypt32.CertAddCertificateContextToStore(valid_store, p_ctx, CERT_STORE_ADD_ALWAYS, None):
                valid_count += 1
        except Exception as exc:
            print(f"   [CertFilter] Bỏ qua chứng thư lỗi: {exc}")

    return valid_store, valid_count

class WindowsStoreSigner(Signer):
    def __init__(self, cert_context_ptr, h_store=None):
        self.cert_context_ptr = cert_context_ptr
        self.h_store = h_store # Giữ store sống để driver không bị unload
        # 1. Tải dữ liệu cert thô và CN trước
        cert_obj, common_name, org_name, subject_dn = self._extract_cert_info()
        self.common_name = common_name
        self.org_name = org_name
        self.subject_dn = subject_dn
        self.cert_data = cert_obj.dump()

        # 2. Lấy toàn bộ chain (Intermediate CA + Root CA) từ Windows Store
        #    → nhúng vào PKCS7 để Adobe / Foxit verify được "Signature Valid"
        self._cert_chain = self._build_cert_chain()

        # 3. Khởi tạo class cha với cert object chuẩn
        super().__init__(signing_cert=cert_obj)

    def _build_cert_chain(self):
        """
        Trích xuất toàn bộ chain từ Windows Certificate Store.
        Trả về list[asn1crypto.x509.Certificate] (không bao gồm chính cert của signer).
        """
        chain = []
        try:
            # Mở thêm các store chứa CA chain
            for store_name in ("CA", "ROOT", "My"):
                h = crypt32.CertOpenSystemStoreW(None, store_name)
                if not h:
                    continue
                # Duyệt tất cả cert trong store
                p = None
                while True:
                    p = crypt32.CertEnumCertificatesInStore(h, p)
                    if not p:
                        break
                    try:
                        from ctypes import Structure, POINTER, c_byte
                        class CERT_CONTEXT(Structure):
                            _fields_ = [
                                ("dwCertEncodingType", wintypes.DWORD),
                                ("pbCertEncoded", POINTER(c_byte)),
                                ("cbCertEncoded", wintypes.DWORD),
                                ("pCertInfo", ctypes.c_void_p),
                                ("hCertStore", ctypes.c_void_p),
                            ]
                        ctx = ctypes.cast(p, POINTER(CERT_CONTEXT)).contents
                        if ctx.pbCertEncoded and ctx.cbCertEncoded > 0:
                            der = ctypes.string_at(ctx.pbCertEncoded, ctx.cbCertEncoded)
                            # Bỏ qua nếu trùng với cert của signer
                            if der != self.cert_data:
                                try:
                                    ca_cert = asn1_x509.Certificate.load(der)
                                    # Chỉ lấy CA cert (không phải end-entity)
                                    bc = ca_cert.ca
                                    chain.append(ca_cert)
                                except Exception:
                                    pass
                    except Exception:
                        pass
                crypt32.CertCloseStore(h, 0)
            print(f"   [Chain] Đã thu thập {len(chain)} CA cert(s) từ Windows Store.")
        except Exception as e:
            print(f"   [Chain] Cảnh báo: Không lấy được chain đầy đủ: {e}")
        return chain

    @property
    def other_certs(self):
        """Pyhanko gọi property này để nhúng các cert phụ vào PKCS7."""
        return self._cert_chain

    def __del__(self):
        # Giải phóng Cert Context khi đối tượng bị hủy
        if hasattr(self, 'cert_context_ptr') and self.cert_context_ptr:
            crypt32.CertFreeCertificateContext(self.cert_context_ptr)
            self.cert_context_ptr = None
        # Đóng store nếu có
        if hasattr(self, 'h_store') and self.h_store:
            crypt32.CertCloseStore(self.h_store, 0)
            self.h_store = None

    def _extract_cert_info(self):
        """Trích xuất đối tượng Certificate và CN từ cert_context_ptr"""
        from ctypes import Structure, POINTER, c_byte
        
        class CERT_CONTEXT(Structure):
            _fields_ = [
                ("dwCertEncodingType", wintypes.DWORD),
                ("pbCertEncoded", POINTER(c_byte)),
                ("cbCertEncoded", wintypes.DWORD),
                ("pCertInfo", ctypes.c_void_p),
                ("hCertStore", ctypes.c_void_p),
            ]
        
        ctx_ptr = ctypes.cast(self.cert_context_ptr, POINTER(CERT_CONTEXT))
        ctx = ctx_ptr.contents
        
        if not ctx.pbCertEncoded:
            raise Exception("pbCertEncoded is NULL")
            
        der_data = ctypes.string_at(ctx.pbCertEncoded, ctx.cbCertEncoded)
        cert = asn1_x509.Certificate.load(der_data)
        subject = cert.subject.native
        
        # Log để chẩn đoán (Diagnostic LOG)
        print(f"   [Debug] Certificate Subject native: {subject}")
        
        # 1. Trích xuất các trường cơ bản
        cn = subject.get('common_name', 'Unknown')
        org = subject.get('organization_name', cn)
        country = subject.get('country_name', 'VN')
        email = subject.get('email_address', '')
        
        # 2. V13.6: Trích xuất Mã Số Thuế (MST) - Quét sâu hơn
        mst = ""
        # Danh sách các trường tiềm năng chứa MST trong Token Việt Nam
        potential_fields = [
            'user_id',              # Token VN thường dùng 'user_id' = 'MST:0311942758'
            'organization_identifier', 
            'serial_number', 
            'userid',               # Một số token dùng không có gạch dưới
            'uid', 
            'title', 
            'description'
        ]
        
        for field in potential_fields:
            val = subject.get(field, '')
            if val:
                val_str = str(val)
                # Nếu value đã có prefix 'MST:' → lấy phần sau
                if val_str.upper().startswith('MST:'):
                    mst = val_str[4:].strip()
                    break
                # MST doanh nghiệp VN thường dài 10 hoặc 13 chữ số
                import re
                match = re.search(r'\d{10,13}', val_str)
                if match:
                    mst = match.group(0)
                    break

        
        # 3. Loại bỏ dấu Tiếng Việt (Accent Stripping) để tránh rác chữ (&D&1...) trên standard fonts
        def strip_accents(s):
            import unicodedata
            if not isinstance(s, str): return str(s)
            s = unicodedata.normalize('NFD', s)
            s = ''.join(c for c in s if unicodedata.category(c) != 'Mn')
            return s.replace('đ', 'd').replace('Đ', 'D')

        self.common_name_clean = strip_accents(cn)
        self.mst = mst
        self.email = email
        self.country = country
        
        # Lưu lại bản gốc để dùng nếu cần
        self.common_name = cn

        # Tự động nhận diện TSA URL từ tên CA trong Issuer
        self.tsa_url = self._detect_tsa_from_issuer(cert)

        print(f"   [Log] Đã trích xuất: {cn} | MST: {mst if mst else 'N/A'} | Email: {email}")
        print(f"   [Log] TSA tự động: {self.tsa_url or 'Không nhận diện được CA VN'}")
        return cert, cn, org, cert.subject.human_friendly

    @staticmethod
    def _detect_tsa_from_issuer(cert) -> str:
        """
        Nhận diện TSA URL từ tên CA trong trường Issuer của certificate.
        Trả về URL TSA tương ứng hoặc chuỗi rỗng nếu không nhận ra CA.
        """
        try:
            issuer = cert.issuer.human_friendly.lower()
        except Exception:
            return ""

        # Bảng tra TSA theo CA Việt Nam (ưu tiên từ trên xuống)
        _TSA_MAP = [
            # Viettel-CA
            (["viettel"],              "http://ca.viettel.vn/tsa"),
            # VNPT-CA / SafeCA
            (["vnpt", "safeca"],       "http://ca.vnpt.vn/tsa"),
            # FPT-CA
            (["fpt"],                  "http://ca.fpt.com.vn/tsa"),
            # BKAV-CA
            (["bkav"],                 "http://ca.bkav.com.vn/tsa"),
            # SmartSign
            (["smartsign", "smart sign"], "http://ca.smartsign.vn/tsa"),
            # MISA
            (["misa"],                 "http://ca.misa.com.vn/tsa"),
            # Nacencomm / Bộ Quốc phòng
            (["nacencomm", "nacentcomm"], "http://ca.nacencomm.vn/tsa"),
            # CK-CA / Chứng khoán
            (["ck-ca", "ckca"],        "http://ca.ckca.vn/tsa"),
            # NEWCA
            (["newca"],                "http://ca.newca.vn/tsa"),
        ]

        for keywords, tsa in _TSA_MAP:
            if any(k in issuer for k in keywords):
                return tsa

        # Không nhận diện được → trả về rỗng (sẽ dùng fallback từ config)
        return ""

    def _load_cert_data(self):
        # Hàm này không còn dùng nữa nhưng giữ lại để tránh lỗi import nếu có
        pass

    async def async_sign_raw(self, data: bytes, digest_algorithm: str, dry_run=False) -> bytes:
        if dry_run:
            return b'\x00' * 256 # Ước tính kích thước chữ ký RSA 2048

        # 0. Hash dữ liệu thô (thông tin SignedAttributes) bằng thuật toán yêu cầu (thường là sha256)
        print(f"   [Log] Đang băm dữ liệu bằng {digest_algorithm}...")
        hasher = hashlib.new(digest_algorithm)
        hasher.update(data)
        hash_value = hasher.digest()
        
        # 1. Thông tin tổ chức đã được trích xuất khi init (self.org_name)

        # 2. Lấy private key handle từ cert context
        h_key = ctypes.c_void_p()
        dw_spec = wintypes.DWORD()
        b_free = wintypes.BOOL()
        
        # Thử nhiều phương án Flag để lách lỗi busy/access (NTE_FAIL / 0x80090020)
        # 0x00000001: CRYPT_ACQUIRE_CACHE_FLAG
        # 0x00000002: CRYPT_ACQUIRE_USE_PROV_INFO_FLAG (Quan trọng cho USB Token VN)
        # 0x00010000: ALLOW_NCRYPT_KEY
        
        import time
        success = False
        last_err = 0
        
        # Danh sách các tổ hợp Flag có độ ổn định giảm dần
        # 0x00000003 = CACHE | USE_PROV_INFO
        flags_to_try = [0x00000003, 0x00000002, 0x00000001, 0x0, 0x00010000]
        
        for attempt in range(4): # Tăng lên 4 lần thử
            for flag in flags_to_try:
                if crypt32.CryptAcquireCertificatePrivateKey(self.cert_context_ptr, flag, None, ctypes.byref(h_key), ctypes.byref(dw_spec), ctypes.byref(b_free)):
                    success = True
                    break
                last_err = ctypes.GetLastError()
            
            if success: break
            
            # --- V12.8: GUI POPUP REMINDER ---
            last_err = last_err & 0xFFFFFFFF # Luôn đưa về số dương (Unsigned)
            if last_err == 0x80090020 or last_err == 0x80090016:
                print(f"   [!] NHẮC NHỞ: Đang chờ USB Token... (Lần {attempt+1}/4)")
                # Hiện Popup và chờ người dùng nhấn OK
                ctypes.windll.user32.MessageBoxW(0, 
                    "Vui lòng gắn USB Token vào máy tính để ký.\n\nSau khi cắm USB, hãy nhấn OK để chương trình tiếp tục.", 
                    "Yêu cầu USB Token", 
                    0x40 | 0x0 # MB_ICONINFORMATION | MB_OK
                )
            else:
                print(f"   [Log] Đang kết nối tới phần cứng (Mã: 0x{last_err:X}). Chờ 1.5s (Lần {attempt+1}/4)...")
                time.sleep(1.5)

        if not success:
            # --- V12.4: VIỆT HÓA CHẾ ĐỘ CỨU HỘ ---
            print(f"   [Log] Tự động kết nối không thành công. Đang kích hoạt Chế độ Cứu hộ (V12.4)...")
            prop_size = wintypes.DWORD()
            if crypt32.CertGetCertificateContextProperty(self.cert_context_ptr, CRYPT_KEY_PROV_INFO_PROP_ID, None, ctypes.byref(prop_size)):
                buf = ctypes.create_string_buffer(prop_size.value)
                if crypt32.CertGetCertificateContextProperty(self.cert_context_ptr, CRYPT_KEY_PROV_INFO_PROP_ID, buf, ctypes.byref(prop_size)):
                    prov_info = ctypes.cast(buf, ctypes.POINTER(CRYPT_KEY_PROV_INFO)).contents
                    
                    if prov_info.dwKeySpec == 0xFFFFFFFF: # Modern CNG
                        h_prov_ncrypt = ctypes.c_void_p()
                        status = ncrypt.NCryptOpenStorageProvider(ctypes.byref(h_prov_ncrypt), prov_info.pwszProvName, 0)
                        if status == 0:
                            status = ncrypt.NCryptOpenKey(h_prov_ncrypt, ctypes.byref(h_key), prov_info.pwszContainerName, 0, 0)
                            if status == 0:
                                dw_spec.value = 0xFFFFFFFF
                                b_free.value = True
                                success = True
                                print(f"   [Log] Manual NCrypt Acquisition thành công.")
                            ncrypt.NCryptFreeObject(h_prov_ncrypt)
                    else: # Legacy CSP
                        h_prov_legacy = ctypes.c_void_p()
                        if advapi32.CryptAcquireContextW(ctypes.byref(h_prov_legacy), prov_info.pwszContainerName, prov_info.pwszProvName, prov_info.dwProvType, 0):
                            h_key.value = h_prov_legacy.value
                            dw_spec.value = prov_info.dwKeySpec
                            b_free.value = True
                            success = True
                            print(f"   [Log] Manual Legacy Acquisition thành công.")

        if not success:
            last_err = last_err & 0xFFFFFFFF # Đảm bảo mã lỗi cuối cùng là số dương
            # Thêm thông tin chẩn đoán
            if last_err == 0x80090020: # NTE_FAIL
                msg = "Lỗi (0x80090020): Vui lòng gắn USB Token vào để ký (Token đang bận hoặc chưa cắm)."
            elif last_err == 0x80090016: # NTE_BAD_KEYSET
                msg = "Lỗi THIẾU THIẾT BỊ (0x80090016): Không tìm thấy Private Key. Vui lòng đảm bảo USB Token của chứng chỉ này đã được cắm vào máy tính."
            else:
                msg = f"Mã lỗi kỹ thuật: 0x{last_err:X}."
            
            raise Exception(msg)

        print(f"   [Log] KeySpec trích xuất được: {dw_spec.value} (0xFFFFFFFF = NCrypt, còn lại = Legacy CSP)")
        

        try:
            if dw_spec.value == 0xFFFFFFFF: # CERT_NCRYPT_KEY_SPEC
                # --- NCrypt Path (CNG) ---
                pad_info = BCRYPT_PKCS1_PADDING_INFO(pszAlgId=digest_algorithm.upper())
                sig_size = wintypes.DWORD()
                print(f"   [Log] Đang dùng NCrypt để ký hash ({len(hash_value)} bytes)...")
                
                status = ncrypt.NCryptSignHash(h_key, ctypes.byref(pad_info), hash_value, len(hash_value), None, 0, ctypes.byref(sig_size), NCRYPT_PAD_PKCS1_FLAG)
                if status != 0:
                    # Fallback: Một số driver không thích pszAlgId
                    print(f"   [Log] NCrypt (with padding info) failed (0x{status:X}). Thử ký không Algorithm ID...")
                    status = ncrypt.NCryptSignHash(h_key, None, hash_value, len(hash_value), None, 0, ctypes.byref(sig_size), NCRYPT_PAD_PKCS1_FLAG)
                
                if status != 0:
                    raise Exception(f"Lỗi NCryptSignHash (get size): 0x{status:X}")

                signature = ctypes.create_string_buffer(sig_size.value)
                status = ncrypt.NCryptSignHash(h_key, ctypes.byref(pad_info), hash_value, len(hash_value), signature, sig_size.value, ctypes.byref(sig_size), NCRYPT_PAD_PKCS1_FLAG)
                if status != 0:
                    # Fallback lần 2
                    status = ncrypt.NCryptSignHash(h_key, None, hash_value, len(hash_value), signature, sig_size.value, ctypes.byref(sig_size), NCRYPT_PAD_PKCS1_FLAG)
                
                if status != 0:
                    raise Exception(f"Lỗi NCryptSignHash (sign): 0x{status:X}")

                print(f"   [Log] Ký (NCrypt) thành công.")
                return bytes(signature)
            else:
                # --- Legacy CSP Path ---
                print(f"   [Log] Đang xử lý Legacy CSP...")
                
                # Sửa lỗi NTE_BAD_ALGID: Nâng cấp handle lên PROV_RSA_AES để hỗ trợ SHA256
                # 1. Lấy thông container name từ cert context
                prop_size = wintypes.DWORD()
                h_key_final = h_key
                b_free_final = b_free
                
                if crypt32.CertGetCertificateContextProperty(self.cert_context_ptr, CRYPT_KEY_PROV_INFO_PROP_ID, None, ctypes.byref(prop_size)):
                    buf = ctypes.create_string_buffer(prop_size.value)
                    if crypt32.CertGetCertificateContextProperty(self.cert_context_ptr, CRYPT_KEY_PROV_INFO_PROP_ID, buf, ctypes.byref(prop_size)):
                        prov_info = ctypes.cast(buf, ctypes.POINTER(CRYPT_KEY_PROV_INFO)).contents
                        
                        # Kỹ thuật quan trọng: Sử dụng Provider AES của Microsoft để "bắc cầu" vào container của Token
                        # Điều này cho phép hỗ trợ SHA-256 ngay cả khi Driver của Token không hỗ trợ.
                        MS_ENH_RSA_AES_PROV = "Microsoft Enhanced RSA and AES Cryptographic Provider"
                        print(f"   [Log] Đang bắc cầu vào Container: {prov_info.pwszContainerName} qua Microsoft AES Provider...")
                        
                        h_prov_aes = ctypes.c_void_p()
                        # Thử mở bằng Provider AES hiện đại của Windows trỏ vào container của Token
                        if advapi32.CryptAcquireContextW(ctypes.byref(h_prov_aes), prov_info.pwszContainerName, MS_ENH_RSA_AES_PROV, PROV_RSA_AES, 0):
                            if b_free.value: advapi32.CryptReleaseContext(h_key, 0)
                            h_key_final = h_prov_aes
                            b_free_final = wintypes.BOOL(True)
                            print(f"   [Log] Bắc cầu (RSA_AES Bridge) thành công.")
                        else:
                            print(f"   [Log] Không thể bắc cầu qua Microsoft AES. Thử dùng Provider gốc: {prov_info.pwszProvName}")
                            # Nếu fail, thử nạp RSA_AES nhưng dùng Prov gốc (đã thử và fail ở bước trước nhưng giữ làm fallback)
                            if advapi32.CryptAcquireContextW(ctypes.byref(h_prov_aes), prov_info.pwszContainerName, prov_info.pwszProvName, PROV_RSA_AES, 0):
                                if b_free.value: advapi32.CryptReleaseContext(h_key, 0)
                                h_key_final = h_prov_aes
                                b_free_final = wintypes.BOOL(True)
                                print(f"   [Log] Nâng cấp Handle bằng Prov gốc thành công.")

                h_hash = ctypes.c_void_p()
                # CALG_SHA256 = 0x800c
                CALG_SHA256 = 0x0000800c
                
                if not advapi32.CryptCreateHash(h_key_final, CALG_SHA256, 0, 0, ctypes.byref(h_hash)):
                    err = ctypes.GetLastError()
                    print(f"   [Error] CryptCreateHash failed (0x{err:X}). Token có thể không hỗ trợ SHA256 qua CSP này.")
                    raise Exception(f"Lỗi CryptCreateHash: 0x{err:X} (NTE_BAD_ALGID=0x80090008?)")
                
                try:
                    # HP_HASHVAL = 2
                    if not advapi32.CryptSetHashParam(h_hash, 2, hash_value, 0):
                        err = ctypes.GetLastError()
                        raise Exception(f"Lỗi CryptSetHashParam: 0x{err:X}")
                    
                    sig_size = wintypes.DWORD()
                    # AT_SIGNATURE = 2, AT_KEYEXCHANGE = 1. CryptoAPI cần KeySpec chuẩn.
                    # Thử sign với dw_spec thu được từ CryptAcquireCertificatePrivateKey
                    if not advapi32.CryptSignHashW(h_hash, dw_spec.value, None, 0, None, ctypes.byref(sig_size)):
                        err = ctypes.GetLastError()
                        raise Exception(f"Lỗi CryptSignHash (get size): 0x{err:X}")
                    
                    signature = ctypes.create_string_buffer(sig_size.value)
                    if not advapi32.CryptSignHashW(h_hash, dw_spec.value, None, 0, signature, ctypes.byref(sig_size)):
                        err = ctypes.GetLastError()
                        raise Exception(f"Lỗi CryptSignHash (sign): 0x{err:X}")
                    
                    print(f"   [Log] Ký (Legacy CSP) thành công.")
                    # Legacy CSP trả về Little-Endian, cần đảo ngược (Reverse) để ra Big-Endian cho PDF
                    return bytes(signature)[::-1]
                finally:
                    advapi32.CryptDestroyHash(h_hash)
                    if b_free_final.value:
                        advapi32.CryptReleaseContext(h_key_final, 0)
                    # Đánh dấu đã free để khối finally ngoài không free lại
                    b_free = wintypes.BOOL(False)

        finally:
            if b_free.value and h_key:
                if dw_spec.value == 0xFFFFFFFF:
                    ncrypt.NCryptFreeObject(h_key)
                else:
                    advapi32.CryptReleaseContext(h_key, 0)

    def get_certificate(self):
        return self.cert_data

def select_signer():
    source_store = crypt32.CertOpenSystemStoreW(None, "My")
    if not source_store:
        return None

    h_store, valid_count = _build_valid_certificate_store(source_store)
    crypt32.CertCloseStore(source_store, 0)

    if not h_store or valid_count <= 0:
        if h_store:
            crypt32.CertCloseStore(h_store, 0)
        try:
            ctypes.windll.user32.MessageBoxW(
                0,
                "Không tìm thấy chứng thư số còn hạn sử dụng trong Windows Store.\n"
                "Vui lòng kiểm tra USB Token hoặc gia hạn chứng thư số.",
                "Chứng thư số",
                0x00000030,
            )
        except Exception:
            pass
        return None

    p_ctx = cryptui.CryptUIDlgSelectCertificateFromStore(h_store, None, "Chọn chứng chỉ ký số", "Chọn chứng chỉ từ USB Token:", 0, 0, None)
    
    if p_ctx:
        print(f"   [Log] Đã chọn chứng chỉ, Context Pointer: {hex(p_ctx)}")
        # Duplicate Certificate Context để sống độc lập
        p_ctx_dup = crypt32.CertDuplicateCertificateContext(p_ctx)
        # Giải phóng context gốc từ dialog
        crypt32.CertFreeCertificateContext(p_ctx)
        # TRẢ VỀ SIGNER KÈM STORE HANDLE (để signer tự đóng khi xong việc)
        return WindowsStoreSigner(p_ctx_dup, h_store)
    else:
        print("   [Log] Người dùng đã hủy chọn chứng chỉ (p_ctx is NULL).")
        crypt32.CertCloseStore(h_store, 0)
    return None


# ===========================================================================
# NÂNG CẤP 1: PKCS#11 Direct Signer
# Ký trực tiếp qua DLL PKCS#11 của USB Token (không qua Windows Store)
# ===========================================================================

class PKCS11Signer(Signer):
    """
    Signer dùng PKCS#11 session để ký trực tiếp từ USB Token.
    Không phụ thuộc vào Windows Certificate Store.
    Tương thích với Viettel-CA, VNPT-CA, FPT-CA, và các CA Việt Nam khác.
    """

    def __init__(self, session, cert_der: bytes, key_obj, label: str = ""):
        """
        Args:
            session:  PyKCS11 session đã login (từ TokenManager)
            cert_der: Certificate ở dạng DER bytes
            key_obj:  Handle private key object từ PKCS#11 session
            label:    Nhãn mô tả của chứng chỉ (tùy chọn)
        """
        from asn1crypto import x509 as asn1_x509
        import unicodedata, re

        self.session = session
        self.cert_der = cert_der
        self.key_obj = key_obj
        self.label = label

        # Trích xuất thông tin từ chứng chỉ
        cert = asn1_x509.Certificate.load(cert_der)
        subject = cert.subject.native
        print(f"   [PKCS11] Certificate Subject: {subject}")

        cn = subject.get('common_name', label or 'Unknown')
        org = subject.get('organization_name', cn)
        email = subject.get('email_address', '')

        # Loại bỏ dấu tiếng Việt
        def strip_accents(s):
            if not isinstance(s, str): return str(s)
            s = unicodedata.normalize('NFD', s)
            s = ''.join(c for c in s if unicodedata.category(c) != 'Mn')
            return s.replace('đ', 'd').replace('Đ', 'D')

        self.common_name = cn
        self.common_name_clean = strip_accents(cn)
        self.org_name = org
        self.email = email

        # Tìm MST
        self.mst = ""
        for field in ['organization_identifier', 'serial_number', 'userid', 'uid']:
            val = subject.get(field, '')
            if val:
                match = re.search(r'\d{10,13}', str(val))
                if match:
                    self.mst = match.group(0)
                    break

        print(f"   [PKCS11] Đã trích xuất: {cn} | MST: {self.mst or 'N/A'} | Email: {email}")

        # Khởi tạo class cha
        super().__init__(signing_cert=cert)

    def get_certificate(self):
        return self.cert_der

    async def async_sign_raw(self, data: bytes, digest_algorithm: str, dry_run=False) -> bytes:
        if dry_run:
            return b'\x00' * 256

        from PyKCS11 import Mechanism
        import asyncio

        # Hash dữ liệu
        print(f"   [PKCS11] Đang băm bằng {digest_algorithm}...")
        hasher = _hashlib.new(digest_algorithm)
        hasher.update(data)
        hash_bytes = hasher.digest()

        # Chọn cơ chế ký phù hợp
        algo_upper = digest_algorithm.upper().replace('-', '')
        mech_map = {
            'SHA256': 0x00000040,   # CKM_SHA256_RSA_PKCS
            'SHA384': 0x00000060,   # CKM_SHA384_RSA_PKCS
            'SHA512': 0x00000041,   # CKM_SHA512_RSA_PKCS
            'SHA1':   0x00000006,   # CKM_SHA1_RSA_PKCS
        }
        ckm_value = mech_map.get(algo_upper, 0x00000040)

        print(f"   [PKCS11] Dùng CKM 0x{ckm_value:08X} để ký hash ({len(hash_bytes)} bytes)...")

        try:
            # Ký raw hash (CKM_RSA_PKCS = 0x00000001)
            mechanism = Mechanism(0x00000001)  # CKM_RSA_PKCS
            loop = asyncio.get_event_loop()
            signature = await loop.run_in_executor(
                None,
                lambda: bytes(self.session.sign(self.key_obj, hash_bytes, mechanism))
            )
            print(f"   [PKCS11] Ký thành công, kích thước: {len(signature)} bytes")
            return signature
        except Exception as e:
            print(f"   [PKCS11] Lỗi khi ký: {e}")
            raise


def select_signer_pkcs11(lib_path: str = None, pin: str = None):
    """
    Mở USB Token qua PKCS#11, cho người dùng chọn chứng chỉ, trả về PKCS11Signer.

    Args:
        lib_path: Đường dẫn đến DLL PKCS#11. Nếu None sẽ tự tìm.
        pin:      PIN của token. Nếu None sẽ hỏi qua console.

    Returns:
        PKCS11Signer hoặc None nếu thất bại.
    """
    import os
    from PyKCS11 import PyKCS11Lib, PyKCS11Error, CKF_SERIAL_SESSION, CKF_RW_SESSION
    from PyKCS11.LowLevel import CKA_CLASS, CKO_CERTIFICATE, CKA_VALUE, CKA_LABEL, CKO_PRIVATE_KEY, CKA_ID
    from asn1crypto import x509 as asn1_x509

    # Tìm DLL tự động
    if not lib_path:
        common_dlls = [
            "viettel-ca_p11.dll", "vnpt-ca_p11.dll", "fpt-ca_p11.dll",
            "bkavca_p11.dll", "ca2_p11.dll", "nacacent_p11.dll",
            "vnpt-ca_v3_p11.dll", "vsign_p11.dll", "dkck201.dll",
            "eps2003csp11.dll", "smartsign_p11.dll"
        ]
        search_dirs = [
            os.environ.get("SystemRoot", r"C:\Windows") + r"\System32",
            os.environ.get("SystemRoot", r"C:\Windows") + r"\SysWOW64",
        ]
        for d in search_dirs:
            for dll in common_dlls:
                full = os.path.join(d, dll)
                if os.path.exists(full):
                    lib_path = full
                    print(f"   [PKCS11] Tự động phát hiện DLL: {lib_path}")
                    break
            if lib_path:
                break

    if not lib_path:
        print("   [PKCS11] Lỗi: Không tìm thấy PKCS#11 DLL. Vui lòng chỉ định lib_path trong sign.yaml")
        return None

    try:
        pkcs11 = PyKCS11Lib()
        pkcs11.load(lib_path)
    except PyKCS11Error as e:
        print(f"   [PKCS11] Lỗi nạp DLL '{lib_path}': {e}")
        return None

    # Liệt kê slots có token
    try:
        slots = pkcs11.getSlotList(tokenPresent=True)
    except PyKCS11Error as e:
        print(f"   [PKCS11] Không lấy được slot list: {e}")
        return None

    if not slots:
        print("   [PKCS11] Không phát hiện USB Token. Vui lòng cắm token vào máy.")
        return None

    # Chọn slot (nếu nhiều hơn 1)
    slot = slots[0]
    if len(slots) > 1:
        print("\n   [PKCS11] Phát hiện nhiều Token:")
        for i, s in enumerate(slots):
            try:
                info = pkcs11.getTokenInfo(s)
                print(f"     [{i+1}] Slot {s}: {info.label.strip()} ({info.manufacturerID.strip()})")
            except Exception:
                print(f"     [{i+1}] Slot {s}")
        try:
            choice = int(input("   Chọn số thứ tự token (Enter để chọn 1): ").strip() or "1") - 1
            slot = slots[max(0, min(choice, len(slots)-1))]
        except ValueError:
            slot = slots[0]

    # Lấy PIN
    if not pin:
        import getpass
        try:
            pin = getpass.getpass("   [PKCS11] Nhập PIN của USB Token: ")
        except Exception:
            pin = input("   [PKCS11] Nhập PIN của USB Token: ")

    # Mở session và login
    try:
        session = pkcs11.openSession(slot, CKF_SERIAL_SESSION | CKF_RW_SESSION)
        session.login(pin)
        print("   [PKCS11] Đăng nhập token thành công.")
    except PyKCS11Error as e:
        print(f"   [PKCS11] Lỗi đăng nhập (PIN sai?): {e}")
        return None

    # Liệt kê chứng chỉ
    try:
        cert_objects = session.findObjects([(CKA_CLASS, CKO_CERTIFICATE)])
    except PyKCS11Error as e:
        print(f"   [PKCS11] Không liệt kê được chứng chỉ: {e}")
        session.logout()
        return None

    if not cert_objects:
        print("   [PKCS11] Không tìm thấy chứng chỉ trên token.")
        session.logout()
        return None

    # Thu thập thông tin chứng chỉ
    cert_list = []
    for obj in cert_objects:
        try:
            attrs = session.getAttributeValue(obj, [CKA_VALUE, CKA_LABEL, CKA_ID])
            cert_der = bytes(attrs[0])
            label = str(attrs[1]) if attrs[1] else ""
            key_id = attrs[2]

            cert = asn1_x509.Certificate.load(cert_der)
            is_valid, reason = _is_asn1_cert_currently_valid(cert)
            if not is_valid:
                print(f"   [PKCS11] Bỏ qua chứng thư hết hạn/chưa hiệu lực ({label}): {reason}")
                continue
            subject = cert.subject.native
            cn = subject.get('common_name', label)
            cert_list.append({'der': cert_der, 'label': label, 'cn': cn, 'key_id': key_id})
        except Exception as e:
            print(f"   [PKCS11] Bỏ qua chứng chỉ lỗi: {e}")

    if not cert_list:
        print("   [PKCS11] Không đọc được chứng chỉ nào.")
        session.logout()
        return None

    # Cho người dùng chọn chứng chỉ nếu có nhiều
    selected = cert_list[0]
    if len(cert_list) > 1:
        print("\n   [PKCS11] Các chứng chỉ trên token:")
        for i, c in enumerate(cert_list):
            print(f"     [{i+1}] {c['cn']} (Label: {c['label']})")
        try:
            choice = int(input("   Chọn chứng chỉ (Enter để chọn 1): ").strip() or "1") - 1
            selected = cert_list[max(0, min(choice, len(cert_list)-1))]
        except ValueError:
            selected = cert_list[0]

    # Tìm private key tương ứng
    try:
        key_objects = session.findObjects([
            (CKA_CLASS, CKO_PRIVATE_KEY),
            (CKA_ID, selected['key_id'])
        ])
        if not key_objects:
            # Fallback: lấy private key đầu tiên
            key_objects = session.findObjects([(CKA_CLASS, CKO_PRIVATE_KEY)])

        if not key_objects:
            print("   [PKCS11] Không tìm thấy private key tương ứng trên token.")
            session.logout()
            return None

        key_obj = key_objects[0]
    except PyKCS11Error as e:
        print(f"   [PKCS11] Lỗi tìm private key: {e}")
        session.logout()
        return None

    print(f"   [PKCS11] Đã chọn chứng chỉ: {selected['cn']}")
    return PKCS11Signer(
        session=session,
        cert_der=selected['der'],
        key_obj=key_obj,
        label=selected['label']
    )
