#include <winsock2.h>
#include <ws2tcpip.h>
#include "lan_channel.h"
#include <bcrypt.h>
#include <limits>

namespace ks::lan
{
namespace
{
void checked(NTSTATUS s){if(s<0) throw std::runtime_error("Input channel authentication failed.");}
void put32(Bytes &b,uint32_t v){for(int n=24;n>=0;n-=8)b.push_back(static_cast<unsigned char>(v>>n));}
uint32_t get32(const unsigned char *p){return (uint32_t(p[0])<<24)|(uint32_t(p[1])<<16)|(uint32_t(p[2])<<8)|p[3];}
Bytes keyBytes(const std::string &key)
{
    if(key.size()!=64||key.find_first_not_of("0123456789abcdef")!=key.npos) throw std::runtime_error("Invalid input key.");
    Bytes bytes;for(size_t i=0;i<key.size();i+=2)bytes.push_back(static_cast<unsigned char>(std::stoul(key.substr(i,2),nullptr,16)));return bytes;
}
bool nonblocking(SOCKET s){u_long value=1;BOOL yes=TRUE;return !ioctlsocket(s,FIONBIO,&value)&&!setsockopt(s,IPPROTO_TCP,TCP_NODELAY,(char*)&yes,sizeof(yes));}
void closeSocket(SOCKET &s){if(s!=INVALID_SOCKET){closesocket(s);s=INVALID_SOCKET;}}
}
struct InputCipher::Impl
{
    BCRYPT_ALG_HANDLE alg{};BCRYPT_KEY_HANDLE key{};uint64_t sequence{};
    explicit Impl(const std::string &secret)
    {
        auto bytes=keyBytes(secret);
        try {
            checked(BCryptOpenAlgorithmProvider(&alg,BCRYPT_AES_ALGORITHM,nullptr,0));
            checked(BCryptSetProperty(alg,BCRYPT_CHAINING_MODE,(PUCHAR)BCRYPT_CHAIN_MODE_GCM,sizeof(BCRYPT_CHAIN_MODE_GCM),0));
            checked(BCryptGenerateSymmetricKey(alg,&key,nullptr,0,bytes.data(),ULONG(bytes.size()),0));
        } catch(...) {if(key)BCryptDestroyKey(key);if(alg)BCryptCloseAlgorithmProvider(alg,0);throw;}
        SecureZeroMemory(bytes.data(),bytes.size());
    }
    ~Impl(){if(key)BCryptDestroyKey(key);if(alg)BCryptCloseAlgorithmProvider(alg,0);}
    Bytes crypt(Bytes payload,Bytes &tag,bool encrypt)
    {
        if(sequence==std::numeric_limits<uint64_t>::max())throw std::runtime_error("Input session exhausted.");
        unsigned char nonce[12]{};for(int i=0;i<8;++i)nonce[11-i]=static_cast<unsigned char>(sequence>>(i*8));
        BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO info;BCRYPT_INIT_AUTH_MODE_INFO(info);
        info.pbNonce=nonce;info.cbNonce=sizeof(nonce);info.pbTag=tag.data();info.cbTag=ULONG(tag.size());
        // Protocol/direction are bound by key derivation; sequence and payload length by GCM.
        Bytes result(payload.size());ULONG written{};
        auto fn=encrypt?BCryptEncrypt:BCryptDecrypt;
        checked(fn(key,payload.data(),ULONG(payload.size()),&info,nullptr,0,result.data(),ULONG(result.size()),&written,0));
        if(written!=payload.size())throw std::runtime_error("Invalid input frame size.");
        ++sequence;return result;
    }
};
InputCipher::InputCipher(const std::string &key):impl(std::make_unique<Impl>(key)){}
InputCipher::~InputCipher()=default;
Bytes InputCipher::encryptPayload(const Bytes &plain)
{
    if(plain.empty()||plain.size()>1024)throw std::runtime_error("Invalid file packet size.");
    Bytes tag(16);auto result=impl->crypt(plain,tag,true);result.insert(result.end(),tag.begin(),tag.end());return result;
}
Bytes InputCipher::decryptPayload(const Bytes &packet)
{
    if(packet.size()<17||packet.size()>1040)throw std::runtime_error("Invalid file packet size.");
    Bytes tag(packet.end()-16,packet.end());return impl->crypt(Bytes(packet.begin(),packet.end()-16),tag,false);
}
Bytes InputCipher::encrypt(const InputFrame &f)
{
    Bytes plain;put32(plain,uint32_t(f.type));put32(plain,f.token);put32(plain,uint32_t(f.a));put32(plain,uint32_t(f.b));put32(plain,uint32_t(f.c));
    Bytes tag(16);auto result=impl->crypt(std::move(plain),tag,true);result.insert(result.end(),tag.begin(),tag.end());return result;
}
InputFrame InputCipher::decrypt(const Bytes &packet)
{
    if(packet.size()!=36)throw std::runtime_error("Invalid encrypted input frame.");
    Bytes tag(packet.begin()+20,packet.end());auto plain=impl->crypt(Bytes(packet.begin(),packet.begin()+20),tag,false);
    auto type=get32(plain.data());if(type>uint32_t(InputType::Cross))throw std::runtime_error("Unsupported input frame.");
    return {InputType(type),get32(plain.data()+4),int32_t(get32(plain.data()+8)),int32_t(get32(plain.data()+12)),int32_t(get32(plain.data()+16))};
}
std::string inputSessionKey(const std::string &secret,const std::string &from,const std::string &to,const std::string &a,const std::string &b,int edge)
{
    if(!identifier(from)||!identifier(to)||from==to||!identifier(a)||!identifier(b))throw std::runtime_error("Invalid session identity.");
    return mac(secret,"KSwitcher input v1|"+from+"|"+to+"|"+a+"|"+b+"|"+std::to_string(edge));
}
struct InputChannel::Impl
{
    std::string local,peer,secret,nonce,state="Waiting for the paired computer...";
    int edge;uint16_t port;SOCKET listener=INVALID_SOCKET,socket=INVALID_SOCKET;
    bool connecting{},authenticated{},helloRead{},managed{};
    uint64_t changed{},lastRead{},lastPing{},retry{};
    Bytes incoming,outgoing;
    std::unique_ptr<InputCipher> tx,rx;
    Impl(std::string a,std::string b,std::string key,int side,uint16_t p,bool multiple):local(std::move(a)),peer(std::move(b)),secret(std::move(key)),edge(side),port(p),managed(multiple){}
    ~Impl(){closeSocket(socket);closeSocket(listener);}
    void reset(uint64_t now,const std::string &reason)
    {
        closeSocket(socket);connecting=authenticated=helloRead=false;incoming.clear();outgoing.clear();tx.reset();rx.reset();++changed;retry=now+1500;state=reason;
    }
    void begin(uint64_t now)
    {
        nonce=randomHex(16);lastRead=lastPing=now;
        auto hello=Json{{"app","KSwitcher-input"},{"version",managed?2:1},{"id",local},{"to",peer},{"nonce",nonce},{"edge",edge}}.dump();
        put32(outgoing,uint32_t(hello.size()));outgoing.insert(outgoing.end(),hello.begin(),hello.end());
        state="Authenticating paired KSwitcher...";
    }
    void poll(uint32_t address,const std::vector<InputFrame> &frames,std::vector<InputFrame> &received,uint64_t now)
    {
        try
        {
            if(socket==INVALID_SOCKET&&now<retry)return;
            if(!managed && listener==INVALID_SOCKET && local>peer)
            {
                listener=::socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);BOOL yes=TRUE;
                if(listener==INVALID_SOCKET||setsockopt(listener,SOL_SOCKET,SO_EXCLUSIVEADDRUSE,(char*)&yes,sizeof(yes))||!nonblocking(listener)){closeSocket(listener);throw std::runtime_error("Cannot open input listener.");}
                sockaddr_in a{};a.sin_family=AF_INET;a.sin_port=htons(port);
                if(bind(listener,(sockaddr*)&a,sizeof(a))||listen(listener,1)) {closeSocket(listener);throw std::runtime_error("TCP input port is unavailable.");}
            }
            if(socket==INVALID_SOCKET && address && now>=retry)
            {
                if(local<peer)
                {
                    socket=::socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);if(socket==INVALID_SOCKET||!nonblocking(socket))throw std::runtime_error("Cannot open input connection.");
                    sockaddr_in a{};a.sin_family=AF_INET;a.sin_port=htons(port);a.sin_addr.s_addr=address;
                    int result=connect(socket,(sockaddr*)&a,sizeof(a));
                    if(result && WSAGetLastError()!=WSAEWOULDBLOCK)throw std::runtime_error("Paired computer is not accepting input connections.");
                    connecting=result!=0;begin(now);
                    if(connecting)state="Connecting to paired PC on TCP 45884...";
                }
                else if(!managed)
                {
                    sockaddr_in a{};int length=sizeof(a);socket=accept(listener,(sockaddr*)&a,&length);
                    if(socket!=INVALID_SOCKET)
                    {
                        if(a.sin_addr.s_addr!=address){closeSocket(socket);return;}
                        if(!nonblocking(socket))throw std::runtime_error("Cannot configure input connection.");
                        begin(now);
                    }
                }
            }
            if(socket==INVALID_SOCKET)return;
            // Give TCP establishment time for SYN retransmission / firewall approval.
            // Keep the short fail-safe only for an already authenticated input session.
            if(now-lastRead>(authenticated?2000:10000))
                throw std::runtime_error(authenticated ? "Input connection lost (no heartbeat for 2 seconds); control returned locally." :
                    connecting ? "Cannot reach TCP 45884. Enable input on both PCs and check the receiving PC firewall." :
                    "TCP connected, but input authentication timed out. Check pairing and input setup on both PCs.");
            if(connecting)
            {
                fd_set writes,errors;FD_ZERO(&writes);FD_ZERO(&errors);FD_SET(socket,&writes);FD_SET(socket,&errors);timeval wait{};
                if(select(0,nullptr,&writes,&errors,&wait)<=0)return;
                int error=0;int size=sizeof(error);
                if(getsockopt(socket,SOL_SOCKET,SO_ERROR,(char*)&error,&size))error=WSAGetLastError();
                if(error)throw std::runtime_error("Cannot connect to TCP 45884 (Windows error "+std::to_string(error)+"). Enable input on both PCs and check the receiving PC firewall.");
                connecting=false;state="TCP connected; authenticating paired KSwitcher...";
            }
            if(authenticated)
            {
                for(auto &f:frames){auto bytes=tx->encrypt(f);outgoing.insert(outgoing.end(),bytes.begin(),bytes.end());}
                if(now-lastPing>=400){auto bytes=tx->encrypt({InputType::Ping});outgoing.insert(outgoing.end(),bytes.begin(),bytes.end());lastPing=now;}
            }
            if(outgoing.size()>32768)throw std::runtime_error("Input queue overflow; control returned locally.");
            if(!outgoing.empty())
            {
                int count=::send(socket,(char*)outgoing.data(),int(outgoing.size()),0);
                if(count>0)outgoing.erase(outgoing.begin(),outgoing.begin()+count);
                else if(count==0||WSAGetLastError()!=WSAEWOULDBLOCK)throw std::runtime_error("Input send failed.");
            }
            unsigned char data[4096];
            for(int batch=0;batch<8;++batch)
            {
                int count=recv(socket,(char*)data,sizeof(data),0);
                if(count==0)throw std::runtime_error("Other KSwitcher disconnected.");
                if(count<0){if(WSAGetLastError()!=WSAEWOULDBLOCK)throw std::runtime_error("Input receive failed.");break;}
                incoming.insert(incoming.end(),data,data+count);
            }
            if(incoming.size()>32768)throw std::runtime_error("Input receive queue overflow.");
            if(!helloRead && incoming.size()>=4)
            {
                size_t length=get32(incoming.data());if(length>512||length<32)throw std::runtime_error("Invalid input handshake.");
                if(incoming.size()<length+4)return;
                auto hello=Json::parse(incoming.begin()+4,incoming.begin()+4+length);
                if(hello.value("app","")!="KSwitcher-input"||hello.value("version",0)!=(managed?2:1)||hello.value("id","")!=peer||hello.value("to","")!=local)
                    throw std::runtime_error("Unexpected input peer.");
                if(!managed&&hello.value("edge",-1)!=(edge^1))throw std::runtime_error("Screen positions must be opposite on the two computers.");
                auto remoteNonce=hello.value("nonce","");
                tx=std::make_unique<InputCipher>(inputSessionKey(secret,local,peer,nonce,remoteNonce,managed?-1:edge));
                rx=std::make_unique<InputCipher>(inputSessionKey(secret,peer,local,remoteNonce,nonce,managed?-1:(edge^1)));
                incoming.erase(incoming.begin(),incoming.begin()+4+length);helloRead=true;
                auto bytes=tx->encrypt({InputType::Ready});outgoing.insert(outgoing.end(),bytes.begin(),bytes.end());
            }
            while(helloRead && incoming.size()>=36)
            {
                auto f=rx->decrypt(Bytes(incoming.begin(),incoming.begin()+36));incoming.erase(incoming.begin(),incoming.begin()+36);lastRead=now;
                if(!authenticated){if(f.type!=InputType::Ready)throw std::runtime_error("Input received before authentication.");authenticated=true;state="Input connected. Move through the configured screen edge.";}
                else if(f.type==InputType::Ready)throw std::runtime_error("Repeated input handshake.");
                else if(f.type!=InputType::Ping)received.push_back(f);
            }
        }
        catch(const std::exception &e){received.clear();reset(now,e.what());}
    }
};
InputChannel::InputChannel(std::string a,std::string b,std::string key,int edge,uint16_t port,bool managed):impl(std::make_unique<Impl>(std::move(a),std::move(b),std::move(key),edge,port,managed)){}
InputChannel::~InputChannel()=default;
void InputChannel::poll(uint32_t ip,const std::vector<InputFrame>&out,std::vector<InputFrame>&in,uint64_t now){impl->poll(ip,out,in,now);}
bool InputChannel::adopt(uintptr_t accepted,uint64_t now)
{
    if(!impl->managed||impl->socket!=INVALID_SOCKET||!nonblocking(SOCKET(accepted)))return false;
    impl->connecting=false;impl->begin(now);impl->socket=SOCKET(accepted);return true;
}
bool InputChannel::ready()const{return impl->authenticated;}
uint64_t InputChannel::generation()const{return impl->changed;}
std::string InputChannel::status()const{return impl->state;}
void InputChannel::disconnect(){impl->reset(GetTickCount64(),"Input disconnected; local control restored.");}
struct InputHub::Impl
{
    struct Link { InputEndpoint endpoint;std::unique_ptr<InputChannel> channel;uint64_t generation{},epoch{};std::string lastState; };
    struct Pending {SOCKET socket;uint32_t address;uint64_t since;};
    std::string local;uint16_t port;SOCKET listener=INVALID_SOCKET;
    std::map<std::string,Link> links;std::vector<Pending> pending;uint64_t epoch{},retry{};
    std::string error;
    Impl(std::string id,uint16_t p):local(std::move(id)),port(p){}
    ~Impl(){closeSocket(listener);for(auto &p:pending)closeSocket(p.socket);}
    void poll(const std::vector<InputEndpoint> &endpoints,const std::vector<RoutedInput> &out,std::vector<RoutedInput> &in,uint64_t now)
    {
        for(auto it=links.begin();it!=links.end();){
            auto match=std::find_if(endpoints.begin(),endpoints.end(),[&](auto &e){return e.id==it->first&&e.key==it->second.endpoint.key&&e.port==it->second.endpoint.port;});
            if(match==endpoints.end())it=links.erase(it);else ++it;
        }
        for(auto &e:endpoints)if(e.id!=local&&!links.contains(e.id)){
            Link link;link.endpoint=e;link.channel=std::make_unique<InputChannel>(local,e.id,e.key,0,e.port,true);link.epoch=++epoch;links.emplace(e.id,std::move(link));
        }
        // Discovery is a hint, not the lifetime of authenticated input sessions.
        // Keep the last usable address across missed announcements. A new address
        // is used only once a channel is down, never to tear down working input.
        for(auto &e:endpoints){if(e.id==local)continue;auto &link=links.at(e.id);
            if(e.address&&!link.channel->ready()&&e.address!=link.endpoint.address){
                link.channel->disconnect();link.endpoint.address=e.address;
            }
        }
        bool listenNeeded=std::any_of(endpoints.begin(),endpoints.end(),[&](auto &e){return e.id<local;});
        if(!listenNeeded){closeSocket(listener);for(auto &p:pending)closeSocket(p.socket);pending.clear();}
        if(listenNeeded&&listener==INVALID_SOCKET&&now>=retry){
            listener=::socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);BOOL yes=TRUE;sockaddr_in a{};a.sin_family=AF_INET;a.sin_port=htons(port);
            if(listener==INVALID_SOCKET||setsockopt(listener,SOL_SOCKET,SO_EXCLUSIVEADDRUSE,(char*)&yes,sizeof(yes))||!nonblocking(listener)||bind(listener,(sockaddr*)&a,sizeof(a))||listen(listener,SOMAXCONN)){
                auto detail="TCP input port "+std::to_string(port)+" is unavailable (Windows error "+std::to_string(WSAGetLastError())+").";
                if(error!=detail)log(L"LAN input: "+wide(detail));
                closeSocket(listener);retry=now+1500;error=std::move(detail);
            }else {error.clear();log(L"LAN input listener ready on TCP "+std::to_wstring(port));}
        }
        if(listener!=INVALID_SOCKET)for(int n=0;n<16&&pending.size()<64;++n){
            sockaddr_in a{};int size=sizeof(a);SOCKET sock=accept(listener,(sockaddr*)&a,&size);if(sock==INVALID_SOCKET)break;
            if(!nonblocking(sock)){closeSocket(sock);continue;}pending.push_back({sock,a.sin_addr.s_addr,now});
        }
        for(auto it=pending.begin();it!=pending.end();){
            bool done=now-it->since>2000;unsigned char bytes[516];int count=recv(it->socket,(char*)bytes,sizeof(bytes),MSG_PEEK);
            if(count==0||(count<0&&WSAGetLastError()!=WSAEWOULDBLOCK))done=true;
            if(!done&&count>=4){
                auto length=get32(bytes);if(length<32||length>512)done=true;
                else if(count>=int(length+4)){
                    done=true;
                    try{auto hello=Json::parse(bytes+4,bytes+4+length);auto id=hello.value("id","");auto link=links.find(id);
                        if(id<local&&hello.value("to","")==local&&hello.value("version",0)==2&&link!=links.end()&&link->second.channel->adopt(uintptr_t(it->socket),now))it->socket=INVALID_SOCKET;
                    }catch(const std::exception &){}
                }
            }
            if(done){closeSocket(it->socket);it=pending.erase(it);}else ++it;
        }
        for(auto &[id,link]:links){
            std::vector<InputFrame> frames,received;for(auto &f:out)if(f.peer==id&&(!f.epoch||f.epoch==link.epoch))enqueueInput(frames,f.frame);
            link.channel->poll(link.endpoint.address,frames,received,now);
            auto state=link.channel->status();if(state!=link.lastState){link.lastState=state;log(L"LAN input peer "+wide(id.substr(0,8))+L": "+wide(state));}
            if(link.generation!=link.channel->generation()){link.generation=link.channel->generation();link.epoch=++epoch;}
            for(auto &f:received)in.push_back({id,f,link.epoch});
        }
    }
};
InputHub::InputHub(std::string id,uint16_t port):impl(std::make_unique<Impl>(std::move(id),port)){}
InputHub::~InputHub()=default;
void InputHub::poll(const std::vector<InputEndpoint> &peers,const std::vector<RoutedInput> &out,std::vector<RoutedInput> &in,uint64_t now){impl->poll(peers,out,in,now);}
std::map<std::string,uint64_t> InputHub::ready()const{std::map<std::string,uint64_t> result;for(auto &[id,l]:impl->links)if(l.channel->ready())result[id]=l.epoch;return result;}
std::string InputHub::status()const{
    auto count=ready().size();if(count)return "Input connected to "+std::to_string(count)+" computer(s). Move through a neighboring screen edge.";
    if(!impl->error.empty())return impl->error;
    for(auto &[id,l]:impl->links)if(!l.channel->status().empty())return l.channel->status();
    return "Waiting for paired computers...";
}
void InputHub::disconnect(){for(auto &[id,l]:impl->links)l.channel->disconnect();}
}