#include "lan_pairing.h"
#include <bcrypt.h>
#include <algorithm>

namespace ks::lan
{
namespace
{
void checked(NTSTATUS s) { if (s < 0) throw std::runtime_error("Windows cryptography failed."); }
std::string hex(const Bytes &data)
{
    std::string out; out.reserve(data.size()*2);
    for (auto b : data) { out += "0123456789abcdef"[b>>4]; out += "0123456789abcdef"[b&15]; }
    return out;
}
std::string proof(const char *kind, const std::string &from, const std::string &to,
                  const std::string &a, const std::string &b = {})
{ return std::string("KSwitcher LAN v1|") + kind + "|" + from + "|" + to + "|" + a + "|" + b; }
}
std::string randomHex(size_t size)
{
    if (size > 64) throw std::runtime_error("Invalid random size.");
    Bytes bytes(size); checked(BCryptGenRandom(nullptr, bytes.data(), ULONG(size), BCRYPT_USE_SYSTEM_PREFERRED_RNG));
    return hex(bytes);
}
bool identifier(const std::string &s)
{ return s.size()==32 && s.find_first_not_of("0123456789abcdef")==s.npos; }
std::string normalizeCode(const std::string &input)
{
    std::string result;
    for (auto c : input)
    {
        if (c=='-' || c==' ') continue;
        if (c>='A' && c<='F') c+= 'a'-'A';
        result+=c;
    }
    if (!identifier(result)) throw std::runtime_error("Use the same generated 32-character pairing code on both computers.");
    return result;
}
std::string mac(const std::string &key, const std::string &text)
{
    BCRYPT_ALG_HANDLE alg{}; BCRYPT_HASH_HANDLE hash{};
    checked(BCryptOpenAlgorithmProvider(&alg, BCRYPT_SHA256_ALGORITHM, nullptr, BCRYPT_ALG_HANDLE_HMAC_FLAG));
    struct Close { BCRYPT_ALG_HANDLE a; BCRYPT_HASH_HANDLE &h; ~Close(){if(h) BCryptDestroyHash(h); BCryptCloseAlgorithmProvider(a,0);} } close{alg,hash};
    checked(BCryptCreateHash(alg,&hash,nullptr,0,(PUCHAR)key.data(),ULONG(key.size()),0));
    checked(BCryptHashData(hash,(PUCHAR)text.data(),ULONG(text.size()),0));
    Bytes out(32); checked(BCryptFinishHash(hash,out.data(),ULONG(out.size()),0)); return hex(out);
}
bool equal(const std::string &a, const std::string &b)
{
    if (a.size()!=b.size()) return false;
    unsigned diff=0; for(size_t i=0;i<a.size();++i) diff |= unsigned(a[i]^b[i]); return diff==0;
}
Json envelope(const std::string &type, const std::string &id)
{ return {{"app","KSwitcher-LAN"},{"version",1},{"type",type},{"id",id}}; }
Pairing::Pairing(std::string a,std::string b,const std::string &key,uint64_t now)
    : local(std::move(a)),peer(std::move(b)),code(normalizeCode(key)),nonce(randomHex(16)),started(now)
{
    if(!identifier(local)||!identifier(peer)||local==peer) throw std::runtime_error("Invalid KSwitcher peer.");
}
bool Pairing::expired(uint64_t now) const { return now-started >= PairLifetime; }
Json Pairing::offer() const
{
    auto j=envelope("offer",local); j["to"]=peer; j["nonce"]=nonce;
    j["proof"]=mac(code,proof("offer",local,peer,nonce)); return j;
}
std::vector<Json> Pairing::receive(const Json &j,uint64_t now)
{
    if(expired(now)||j.value("app","")!="KSwitcher-LAN"||j.value("version",0)!=1||
       j.value("id","")!=peer||j.value("to","")!=local) return {};
    auto type=j.value("type",""); auto other=j.value("nonce","");
    if(!identifier(other)) return {};
    if(type=="offer")
    {
        if(!equal(j.value("proof",""),mac(code,proof("offer",peer,local,other)))) return {};
        // A new attempt cannot replace a nonce already authenticated in this attempt.
        if(!remoteNonce.empty() && remoteNonce!=other) return {};
        remoteNonce=other;
        derived=local<peer ? mac(code,proof("key",local,peer,nonce,other)) : mac(code,proof("key",peer,local,other,nonce));
        auto ack=envelope("ack",local); ack["to"]=peer; ack["nonce"]=nonce; ack["echo"]=other;
        ack["proof"]=mac(code,proof("ack",local,peer,other,nonce)); return {ack};
    }
    if(type=="ack" && j.value("echo","")==nonce)
    {
        if(!equal(j.value("proof",""),mac(code,proof("ack",peer,local,nonce,other)))) return {};
        if(!remoteNonce.empty()&&remoteNonce!=other) return {};
        remoteNonce=other;
        derived=local<peer ? mac(code,proof("key",local,peer,nonce,other)) : mac(code,proof("key",peer,local,other,nonce));
        acknowledged=true;
        auto done=envelope("confirm",local); done["to"]=peer; done["nonce"]=nonce;
        done["proof"]=mac(derived,proof("confirm",local,peer,nonce,other)); return {done};
    }
    if(type=="confirm" && acknowledged && other==remoteNonce &&
       equal(j.value("proof",""),mac(derived,proof("confirm",peer,local,other,nonce)))) confirmed=true;
    return {};
}
}