import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.keyinfo.*;
import javax.xml.crypto.dsig.spec.*;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.cert.X509Certificate;
import java.time.ZoneOffset;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.*;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import org.w3c.dom.*;

public class KSeFSessionIniter {

    private static final String BASE_URL = "https://api-demo.ksef.mf.gov.pl/v2";

    public static void main(String[] args) {
        try {
            if (args.length < 3) {
                System.out.println("Użycie: java KSeFSessionIniter <nip> <p12_path> <p12_pass>");
                return;
            }

            String nip = args[0];
            String p12Path = args[1];
            String password = args[2];

            KeyStore ks = KeyStore.getInstance("PKCS12");
            ks.load(new FileInputStream(p12Path), password.toCharArray());
            String alias = ks.aliases().nextElement();
            PrivateKey privateKey = (PrivateKey) ks.getKey(alias, password.toCharArray());
            X509Certificate cert = (X509Certificate) ks.getCertificate(alias);

            System.out.println("Pobieranie challenge...");
            String challenge = fetchChallenge();
            
            System.out.println("Generowanie i podpisywanie XML...");
            String signedXml = generateSignedXml(challenge, nip, privateKey, cert);

            System.out.println("Inicjalizacja sesji...");
            String response = sendRequest(BASE_URL + "/auth/xades-signature", "POST", signedXml);
            System.out.println("Odpowiedź serwera:\n" + response);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static String fetchChallenge() throws Exception {
        String url = BASE_URL + "/auth/challenge";
        String response = sendRequest(url, "POST", "{}"); 
        return response.split("\"challenge\":\"")[1].split("\"")[0];
    }

    private static String generateSignedXml(String challenge, String nip, PrivateKey privateKey, X509Certificate cert) throws Exception {
        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);
        Document doc = dbf.newDocumentBuilder().newDocument();

        String ns = "http://ksef.mf.gov.pl/auth/token/2.0";
        Element root = doc.createElementNS(ns, "AuthTokenRequest");
        root.setAttribute("Id", "AuthTokenRequest");
        doc.appendChild(root);

        Element chalElem = doc.createElementNS(ns, "Challenge");
        chalElem.setTextContent(challenge);
        root.appendChild(chalElem);

        Element ident = doc.createElementNS(ns, "ContextIdentifier");
        Element nipElem = doc.createElementNS(ns, "Nip");
        nipElem.setTextContent(nip);
        ident.appendChild(nipElem);
        root.appendChild(ident);

        Element subType = doc.createElementNS(ns, "SubjectIdentifierType");
        subType.setTextContent("certificateSubject");
        root.appendChild(subType);

        XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");

        // Referencja 1: Główny dokument (używamy EXCLUSIVE C14N)
        Reference refDoc = fac.newReference("#AuthTokenRequest", fac.newDigestMethod(DigestMethod.SHA256, null),
                Collections.singletonList(fac.newTransform(Transform.ENVELOPED, (TransformParameterSpec) null)), null, null);

        // Struktura XAdES
        String xadesNs = "http://uri.etsi.org/01903/v1.3.2#";
        String dsigNs = "http://www.w3.org/2000/09/xmldsig#";
        
        Element qualProps = doc.createElementNS(xadesNs, "xades:QualifyingProperties");
        qualProps.setAttribute("Target", "#AuthTokenRequest");
        Element signedProps = doc.createElementNS(xadesNs, "xades:SignedProperties");
        signedProps.setAttribute("Id", "SignedProperties");
        
        Element ssp = doc.createElementNS(xadesNs, "xades:SignedSignatureProperties");
        Element st = doc.createElementNS(xadesNs, "xades:SigningTime");
        st.setTextContent(ZonedDateTime.now(ZoneOffset.UTC).format(DateTimeFormatter.ISO_INSTANT));
        ssp.appendChild(st);
        
        Element sc = doc.createElementNS(xadesNs, "xades:SigningCertificate");
        Element c = doc.createElementNS(xadesNs, "xades:Cert");
        
        Element cd = doc.createElementNS(xadesNs, "xades:CertDigest");
        Element dm = doc.createElementNS(dsigNs, "ds:DigestMethod");
        dm.setAttribute("Algorithm", "http://www.w3.org/2001/04/xmlenc#sha256");
        MessageDigest md = MessageDigest.getInstance("SHA-256");
        byte[] certHash = md.digest(cert.getEncoded());
        Element dv = doc.createElementNS(dsigNs, "ds:DigestValue");
        dv.setTextContent(Base64.getEncoder().encodeToString(certHash));
        cd.appendChild(dm); cd.appendChild(dv); c.appendChild(cd);

        Element iser = doc.createElementNS(xadesNs, "xades:IssuerSerial");
        Element iname = doc.createElementNS(dsigNs, "ds:X509IssuerName");
        iname.setTextContent(cert.getIssuerX500Principal().getName());
        Element iserNum = doc.createElementNS(dsigNs, "ds:X509SerialNumber");
        iserNum.setTextContent(cert.getSerialNumber().toString());
        iser.appendChild(iname); iser.appendChild(iserNum);
        c.appendChild(iser);
        
        sc.appendChild(c); ssp.appendChild(sc); signedProps.appendChild(ssp); qualProps.appendChild(signedProps);

        // Referencja 2: SignedProperties (używamy EXCLUSIVE C14N)
        Reference refXades = fac.newReference("#SignedProperties", fac.newDigestMethod(DigestMethod.SHA256, null),
                Collections.singletonList(fac.newTransform(CanonicalizationMethod.EXCLUSIVE, (C14NMethodParameterSpec) null)),
                "http://uri.etsi.org/01903/v1.3.2#SignedProperties", null);

        String sigMethod = privateKey.getAlgorithm().equals("EC") 
            ? "http://www.w3.org/2001/04/xmldsig-more#ecdsa-sha256" 
            : SignatureMethod.RSA_SHA256;

        SignedInfo si = fac.newSignedInfo(fac.newCanonicalizationMethod(CanonicalizationMethod.EXCLUSIVE, (C14NMethodParameterSpec) null),
                fac.newSignatureMethod(sigMethod, null), Arrays.asList(refDoc, refXades));

        KeyInfoFactory kif = fac.getKeyInfoFactory();
        X509Data xd = kif.newX509Data(Collections.singletonList(cert));
        KeyInfo ki = kif.newKeyInfo(Collections.singletonList(xd));

        XMLObject xadesObj = fac.newXMLObject(Collections.singletonList(new javax.xml.crypto.dom.DOMStructure(qualProps)), null, null, null);
        XMLSignature signature = fac.newXMLSignature(si, ki, Collections.singletonList(xadesObj), null, null);

        DOMSignContext dsc = new DOMSignContext(privateKey, root);
        dsc.setIdAttributeNS(root, null, "Id");
        dsc.setIdAttributeNS(signedProps, null, "Id");

        doc.normalize();
        signature.sign(dsc);

        Transformer trans = TransformerFactory.newInstance().newTransformer();
        trans.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "yes");
        trans.setOutputProperty(OutputKeys.INDENT, "no");

        StringWriter sw = new StringWriter();
        trans.transform(new DOMSource(doc), new StreamResult(sw));
        
        return sw.toString().replaceAll(">\\s+<", "><").trim();
    }

    private static String sendRequest(String urlStr, String method, String body) throws Exception {
        URL url = new URL(urlStr);
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod(method);
        conn.setRequestProperty("Accept", "application/json");
        conn.setRequestProperty("Content-Type", urlStr.contains("xades") ? "application/xml" : "application/json");
        conn.setDoOutput(true);

        if (body != null) {
            try (OutputStream os = conn.getOutputStream()) {
                os.write(body.getBytes(StandardCharsets.UTF_8));
            }
        }

        int status = conn.getResponseCode();
        InputStream is = (status < 400) ? conn.getInputStream() : conn.getErrorStream();
        if (is == null) return "{\"error\": \"HTTP " + status + "\"}";

        try (BufferedReader br = new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) {
            StringBuilder sb = new StringBuilder();
            String line;
            while ((line = br.readLine()) != null) sb.append(line);
            return sb.toString();
        }
    }
}