#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import os
import json
import base64
from typing import List, Dict, Optional, Tuple

import requests
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding as sym_padding

# ------------- helpers low-level crypto -------------

def gen_aes_key_iv() -> Tuple[bytes, bytes]:
    key = os.urandom(32)  # 256-bit
    iv  = os.urandom(16)  # 128-bit IV
    return key, iv

def aes256_cbc_pkcs7_encrypt(plaintext: bytes, key: bytes, iv: bytes) -> bytes:
    padder = sym_padding.PKCS7(128).padder()
    padded = padder.update(plaintext) + padder.finalize()

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    enc = cipher.encryptor()
    ciphertext = enc.update(padded) + enc.finalize()
    return ciphertext

def b64_sha256(data: bytes) -> str:
    h = hashes.Hash(hashes.SHA256())
    h.update(data)
    digest = h.finalize()
    return base64.b64encode(digest).decode("ascii")

def rsa_oaep_sha256_encrypt(data: bytes, cert: x509.Certificate) -> bytes:
    pub = cert.public_key()
    return pub.encrypt(
        data,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )

# ------------- KSeF API helpers -------------

def get_public_key_certificates(base_url: str) -> List[Dict]:
    url = base_url.rstrip("/") + "/api/v2/security/public-key-certificates"
    r = requests.get(url, headers={"Accept": "application/json"}, timeout=30)
    r.raise_for_status()
    js = r.json()
    if not isinstance(js, list):
        raise RuntimeError("Nieoczekiwana odpowiedź /public-key-certificates")
    return js

def der_b64_to_cert(der_b64: str) -> x509.Certificate:
    der = base64.b64decode(der_b64)
    return x509.load_der_x509_certificate(der)

def pick_cert_for_usage(items: List[Dict], wanted_usage: str) -> x509.Certificate:
    """
    wanted_usage np. 'SymmetricKeyEncryption'
    """
    for item in items:
        usage = item.get("usage") or []
        if wanted_usage in usage:
            return der_b64_to_cert(item["certificate"])
    raise RuntimeError(f"Brak certyfikatu z usage={wanted_usage}")

def open_online_session(base_url: str,
                        access_token: str,
                        enc_sym_key_b64: str,
                        iv_b64: str) -> Dict:
    """
    POST /api/v2/sessions/online

    Minimalny payload:
    {
      "formCode": {
         "systemCode": "FA (3)",
         "schemaVersion": "1-0E",
         "value": "FA"
      },
      "encryption": {
         "encryptedSymmetricKey": "<Base64>",
         "initializationVector": "<Base64>"
      }
    }

    NOTE:
    Jeśli u Ciebie inne wartości formCode (np. "FA", "FA (2)", wersja schematu),
    zmień tutaj, albo podaj to później jako param przez argv.
    Na razie twardo wpisujemy typową konfigurację testową.
    """
    url = base_url.rstrip("/") + "/api/v2/sessions/online"
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {access_token}",
    }
    payload = {
        "formCode": {
            "systemCode": "FA (3)",
            "schemaVersion": "1-0E",
            "value": "FA"
        },
        "encryption": {
            "encryptedSymmetricKey": enc_sym_key_b64,
            "initializationVector": iv_b64
        }
    }

    r = requests.post(url, headers=headers, json=payload, timeout=30)

    out = {
        "status": r.status_code,
        "raw": None,
        "json": None,
    }

    try:
        out["json"] = r.json()
    except Exception:
        out["raw"] = r.text

    return out

def send_online_invoice(base_url: str,
                        access_token: str,
                        session_ref: str,
                        invoice_hash_b64: str,
                        invoice_size: int,
                        enc_hash_b64: str,
                        enc_size: int,
                        enc_content_b64: str,
                        offline_mode: bool = False) -> Dict:
    """
    POST /api/v2/sessions/online/{referenceNumber}/invoices
    {
      "invoiceHash": "<Base64(SHA-256 oryginału)>",
      "invoiceSize": <bytes>,
      "encryptedInvoiceHash": "<Base64(SHA-256 szyfrogramu)>",
      "encryptedInvoiceSize": <bytes>,
      "encryptedInvoiceContent": "<Base64(ciphertext AES-256-CBC)>",
      "offlineMode": false
    }
    """
    url = base_url.rstrip("/") + f"/api/v2/sessions/online/{session_ref}/invoices"
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {access_token}",
    }
    payload = {
        "invoiceHash":            invoice_hash_b64,
        "invoiceSize":            invoice_size,
        "encryptedInvoiceHash":   enc_hash_b64,
        "encryptedInvoiceSize":   enc_size,
        "encryptedInvoiceContent": enc_content_b64,
        "offlineMode":            offline_mode,
    }

    r = requests.post(url, headers=headers, json=payload, timeout=60)

    out = {
        "status": r.status_code,
        "raw": None,
        "json": None,
    }
    try:
        out["json"] = r.json()
    except Exception:
        out["raw"] = r.text

    return out

def close_online_session(base_url: str,
                         access_token: str,
                         session_ref: str) -> Dict:
    """
    POST /api/v2/sessions/online/{referenceNumber}/close
    """
    url = base_url.rstrip("/") + f"/api/v2/sessions/online/{session_ref}/close"
    headers = {
        "Accept": "application/json",
        "Authorization": f"Bearer {access_token}",
    }

    r = requests.post(url, headers=headers, timeout=20)

    out = {
        "status": r.status_code,
        "raw": None,
        "json": None,
    }

    if r.status_code == 204:
        # 204 = brak treści, czyli sukces bez JSON-a
        return {
            "status": 204,
            "json": None,
            "raw": "",
        }

    try:
        out["json"] = r.json()
    except Exception:
        out["raw"] = r.text

    return out


# ------------- main flow -------------

def main():
    # argv:
    # 1: BASE_URL (np. https://ksef-test.mf.gov.pl)
    # 2: accessToken (z redeemTokens, JWT)
    # 3: invoiceFullPath (pełna ścieżka pliku XML faktury)
    if len(sys.argv) < 4:
        raise RuntimeError("Użycie: send_invoice.py BASE_URL ACCESS_TOKEN /full/path/to/faktura.xml")

    base_url   = sys.argv[1].strip()
    access_tok = sys.argv[2].strip()
    invoice_fp = sys.argv[3].strip()

    if not os.path.isfile(invoice_fp):
        raise RuntimeError(f"Brak pliku faktury: {invoice_fp}")

    # 1) pobierz publiczne certyfikaty KSeF
    cert_items = get_public_key_certificates(base_url)
    # wybierz cert do zaszyfrowania klucza symetrycznego
    cert_for_sym = pick_cert_for_usage(cert_items, "SymmetricKeyEncryption")

    # 2) generuj klucz AES 256 + IV
    aes_key, aes_iv = gen_aes_key_iv()

    # 3) zaszyfruj ten klucz publicznym kluczem MF (RSA-OAEP SHA-256)
    enc_key_bin = rsa_oaep_sha256_encrypt(aes_key, cert_for_sym)
    enc_key_b64 = base64.b64encode(enc_key_bin).decode("ascii")
    iv_b64      = base64.b64encode(aes_iv).decode("ascii")

    # 4) otwórz sesję online
    open_resp = open_online_session(base_url, access_tok, enc_key_b64, iv_b64)
    status_open = open_resp.get("status")
    open_json   = open_resp.get("json") or {}
    session_ref = (
        open_json.get("referenceNumber")
        or open_json.get("sessionReferenceNumber")
        or open_json.get("reference")
        or ""
    )
    if status_open not in (200,201) or not session_ref:
        # sesja się nie otworzyła poprawnie
        raise RuntimeError(f"Nie udało się otworzyć sesji online ({status_open}): {open_resp}")

    # 5) wczytaj fakturę i zaszyfruj
    invoice_bytes = open(invoice_fp, "rb").read()
    encrypted_invoice = aes256_cbc_pkcs7_encrypt(invoice_bytes, aes_key, aes_iv)

    # 6) policz metadane (hash oryginału i hash szyfrogramu)
    inv_hash_b64    = b64_sha256(invoice_bytes)
    inv_size        = len(invoice_bytes)
    enc_hash_b64    = b64_sha256(encrypted_invoice)
    enc_size        = len(encrypted_invoice)
    enc_content_b64 = base64.b64encode(encrypted_invoice).decode("ascii")

    # 7) wyślij fakturę
    send_resp = send_online_invoice(
        base_url,
        access_tok,
        session_ref,
        inv_hash_b64,
        inv_size,
        enc_hash_b64,
        enc_size,
        enc_content_b64,
        offline_mode=False
    )
    status_send = send_resp.get("status")

    # 8) zamknij sesję (nawet jeśli wysyłka dała np. 202)
    close_resp   = close_online_session(base_url, access_tok, session_ref)
    status_close = close_resp.get("status")

    # 9) przygotuj wynik dla Laravela
    result = {
        "sessionReferenceNumber": session_ref,
        "statusOpen":  status_open,
        "statusSend":  status_send,
        "statusClose": status_close,
        "detailsOpen": open_resp,
        "detailsSend": send_resp,
        "detailsClose": close_resp,
    }

    # drukujemy JEDEN JSON na STDOUT
    print(json.dumps(result, ensure_ascii=False))

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        # jeżeli cokolwiek pójdzie źle, wypisz błąd na STDERR
        # i wyjdź z kodem 1 -> Laravel to złapie i pokaże w błędzie
        sys.stderr.write(str(e) + "\n")
        sys.exit(1)