"""Motor PAdES local del xamiSDK (Fase 3).

Firma diferida (two-phase) con pyHanko: el documento NUNCA sale de la infra del
cliente. Solo el hash del ByteRange viaja a Xami; el chip firma; el SDK incrusta la
firma localmente.

Flujo:
  prepare(pdf_bytes, credential, design) -> {request_id, digest, state_path}
      calcula el ByteRange y su hash SHA-256, escribe el PDF con un placeholder de
      firma y guarda el estado (PDF preparado + offsets) para el assemble posterior.
  assemble(request_id, signature, cert)  -> pdf_bytes firmado
      construye el CMS (SignedData) con la firma del chip y el cert, y lo incrusta en
      el placeholder. Devuelve el PDF firmado.

Ver docs/servicios_xamiSDK.md (§5.1) y docs/arquitectura_xamiSDK.md.
"""
from __future__ import annotations
import hashlib
import io
import json
import os
import stat
import uuid

from pyhanko.pdf_utils.incremental_writer import IncrementalPdfFileWriter
from pyhanko.sign import fields
from pyhanko.sign.fields import SigFieldSpec
from pyhanko.sign.signers.cms_embedder import PdfCMSEmbedder, SigObjSetup, SigIOSetup
from pyhanko.sign.signers.pdf_signer import SignatureObject
from pyhanko.sign.fields import SigSeedSubFilter


class PadesEngine:
    def __init__(self, work_dir: str):
        # estado intermedio entre prepare y assemble (por request_id)
        self.state_dir = os.path.join(work_dir, "pades_state")
        os.makedirs(self.state_dir, exist_ok=True)

    # ------------------------------------------------------------------ prepare
    def prepare(self, pdf_bytes: bytes, *, field_name: str = "Signature1",
                reason: str | None = None, location: str | None = None,
                signer_name: str | None = None,
                design: dict | None = None, variables: dict | None = None,
                image_path: str | None = None) -> dict:
        """Fase 1: calcula el hash del ByteRange y deja el PDF con placeholder listo.
        Si se da `design` (visible) con su box/page, aplica la apariencia (sello).
        Devuelve {request_id, digest (hex), state_path}."""
        request_id = "sig_" + uuid.uuid4().hex[:16]

        # resolver apariencia desde el diseno en cache (si es visible)
        params = (design or {}).get("params") or {}
        on_page = int(params.get("page", 1)) - 1
        if on_page < 0:
            on_page = 0
        box_spec = self._box_to_pyhanko(params.get("box"))

        # el campo de firma: con box si hay apariencia visible, invisible si no
        if box_spec is not None:
            field_spec = SigFieldSpec(sig_field_name=field_name, on_page=on_page, box=box_spec)
        else:
            field_spec = SigFieldSpec(sig_field_name=field_name, on_page=0)

        writer = IncrementalPdfFileWriter(io.BytesIO(pdf_bytes))
        cms_writer = PdfCMSEmbedder(new_field_spec=field_spec).write_cms(
            field_name=field_name, writer=writer, existing_fields_only=False)

        # fase 1: crear/obtener el campo
        next(cms_writer)

        # apariencia (sello visual) si el diseno es visible
        appearance_setup = None
        if params.get("visible") and box_spec is not None:
            try:
                from xamisdk import design_stamp
                from datetime import datetime as _dt
                style = design_stamp.build_stamp_style(
                    params, image_full_path=image_path, variables=variables or {})
                if style is not None:
                    from pyhanko.sign.signers.cms_embedder import SigAppearanceSetup
                    appearance_setup = SigAppearanceSetup(
                        style=style, timestamp=_dt.now(), name=signer_name or "")
            except Exception:
                appearance_setup = None   # si falla la apariencia, firma sin sello visible

        # fase 2: enviar SigObjSetup con el placeholder (reserva de bytes para el CMS)
        placeholder = SignatureObject(
            subfilter=SigSeedSubFilter.ADOBE_PKCS7_DETACHED,
            name=signer_name, location=location, reason=reason,
            bytes_reserved=8192,
        )
        sig_obj_setup = SigObjSetup(sig_placeholder=placeholder,
                                    appearance_setup=appearance_setup)
        cms_writer.send(sig_obj_setup)

        # fase 3: enviar SigIOSetup -> escribe con placeholder y devuelve el hash
        sig_io = SigIOSetup(md_algorithm="sha256", in_place=False)
        prepared_digest, output = cms_writer.send(sig_io)

        digest_bytes = prepared_digest.document_digest
        digest_hex = digest_bytes.hex()

        # guardar el estado: PDF con placeholder + offsets de la region reservada
        output.seek(0)
        prepared_pdf = output.read()
        state = {
            "request_id": request_id,
            "field_name": field_name,
            "digest": digest_hex,
            "reserved_region_start": prepared_digest.reserved_region_start,
            "reserved_region_end": prepared_digest.reserved_region_end,
        }
        state_path = self._save_state(request_id, state, prepared_pdf)

        return {"request_id": request_id, "digest": digest_hex, "state_path": state_path}

    # ----------------------------------------------------------------- assemble
    def assemble(self, request_id: str, signature_hex: str, cert_pem: str,
                 algorithm: str | None = None) -> bytes:
        """Fase 2: construye el CMS con la firma del chip + cert y lo incrusta en el
        placeholder. Devuelve el PDF firmado. `algorithm` viene del chip
        (ECDSA-P256-... o RSA-...); si no se da, se infiere del cert."""
        state, prepared_pdf = self._load_state(request_id)
        if state is None:
            raise ValueError(f"no hay estado para request_id {request_id}")

        cms_der = self._build_cms(state["digest"], signature_hex, cert_pem, algorithm)

        # incrustar el CMS en la region reservada usando el helper de pyHanko
        import io as _io
        from pyhanko.sign.signers.pdf_byterange import PreparedByteRangeDigest
        prepared = PreparedByteRangeDigest(
            document_digest=bytes.fromhex(state["digest"]),
            reserved_region_start=state["reserved_region_start"],
            reserved_region_end=state["reserved_region_end"],
        )
        out = _io.BytesIO(prepared_pdf)
        prepared.fill_with_cms(out, cms_der)
        out.seek(0)
        return out.read()

    # ------------------------------------------------------------------ helpers
    def _box_to_pyhanko(self, box):
        """Convierte el box del diseno a la tupla (x1,y1,x2,y2) que pyHanko espera para
        el rect del campo. Acepta {x,y,width,height} (objeto) o [x1,y1,x2,y2] (lista).
        Devuelve None si no hay box (firma invisible)."""
        if not box:
            return None
        if isinstance(box, dict):
            try:
                x = float(box["x"]); y = float(box["y"])
                w = float(box["width"]); h = float(box["height"])
                return (x, y, x + w, y + h)
            except Exception:
                return None
        if isinstance(box, (list, tuple)) and len(box) == 4:
            try:
                return tuple(float(v) for v in box)
            except Exception:
                return None
        return None

    def _build_cms(self, digest_hex: str, signature_hex: str, cert_pem: str,
                   algorithm: str | None = None) -> bytes:
        """Construye un SignedData CMS/PKCS#7 detached con la firma del chip.
        Soporta RSA (sha256_rsa) y ECDSA P-256 (sha256_ecdsa), inferido del
        algorithm reportado por el chip o del tipo de clave del cert."""
        from asn1crypto import cms, algos, core, x509 as asn1x509
        from cryptography import x509 as cx509
        from cryptography.hazmat.primitives.serialization import Encoding
        from cryptography.hazmat.primitives.asymmetric import rsa as _rsa, ec as _ec

        cert = cx509.load_pem_x509_certificate(cert_pem.encode())
        cert_der = cert.public_bytes(Encoding.DER)
        asn1_cert = asn1x509.Certificate.load(cert_der)

        # elegir el algoritmo de firma para el signerInfo
        algo_up = (algorithm or "").upper()
        pub = cert.public_key()
        if "RSA" in algo_up or isinstance(pub, _rsa.RSAPublicKey):
            sig_algo = "sha256_rsa"
        elif "ECDSA" in algo_up or "EC " in algo_up or isinstance(pub, _ec.EllipticCurvePublicKey):
            sig_algo = "sha256_ecdsa"
        else:
            sig_algo = "sha256_rsa"   # por defecto PAdES suele ser RSA

        digest_bytes = bytes.fromhex(digest_hex)
        signature_bytes = bytes.fromhex(signature_hex)

        sid = cms.SignerIdentifier({
            "issuer_and_serial_number": cms.IssuerAndSerialNumber({
                "issuer": asn1_cert.issuer,
                "serial_number": asn1_cert.serial_number,
            })
        })

        # Sin signed_attrs: el chip firmo el digest del ByteRange directamente, asi que
        # la firma del signerInfo es sobre ese digest (messageDigest). Esto hace que la
        # firma del chip sea directamente valida en el CMS sin re-firmar atributos.
        signer_info = cms.SignerInfo({
            "version": "v1",
            "sid": sid,
            "digest_algorithm": algos.DigestAlgorithm({"algorithm": "sha256"}),
            "signature_algorithm": algos.SignedDigestAlgorithm({"algorithm": sig_algo}),
            "signature": signature_bytes,
        })

        signed_data = cms.SignedData({
            "version": "v1",
            "digest_algorithms": [algos.DigestAlgorithm({"algorithm": "sha256"})],
            "encap_content_info": {"content_type": "data"},
            "certificates": [asn1_cert],
            "signer_infos": [signer_info],
        })

        content_info = cms.ContentInfo({
            "content_type": "signed_data",
            "content": signed_data,
        })
        return content_info.dump()

    def _save_state(self, request_id: str, state: dict, prepared_pdf: bytes) -> str:
        d = os.path.join(self.state_dir, request_id)
        os.makedirs(d, exist_ok=True)
        with open(os.path.join(d, "state.json"), "w") as f:
            json.dump(state, f)
        pdf_path = os.path.join(d, "prepared.pdf")
        with open(pdf_path, "wb") as f:
            f.write(prepared_pdf)
        os.chmod(pdf_path, stat.S_IRUSR | stat.S_IWUSR)
        return d

    def _load_state(self, request_id: str):
        d = os.path.join(self.state_dir, request_id)
        sp = os.path.join(d, "state.json")
        pp = os.path.join(d, "prepared.pdf")
        if not os.path.exists(sp) or not os.path.exists(pp):
            return None, None
        with open(sp) as f:
            state = json.load(f)
        with open(pp, "rb") as f:
            prepared_pdf = f.read()
        return state, prepared_pdf

    def cleanup(self, request_id: str):
        """Borra el estado intermedio tras ensamblar."""
        import shutil
        d = os.path.join(self.state_dir, request_id)
        if os.path.isdir(d):
            shutil.rmtree(d)
