#include <winsock2.h>
#include <ws2tcpip.h>
#include <iphlpapi.h>
#include "lan.h"
#include "lan_pairing.h"
#include "lan_mesh.h"
#include "lan_clipboard.h"
#include "lan_layout.h"
#include <atomic>
#include <mutex>
#include <thread>
#include <map>

namespace ks
{
using lan::Json;
namespace
{
Bytes encoded(const Json &j) { auto s=j.dump(); return Bytes(s.begin(),s.end()); }
bool localAddress(ULONG network)
{
    auto ip=ntohl(network);
    return (ip>>24)==10 || (ip>>20)==0xac1 || (ip>>16)==0xc0a8 || (ip>>16)==0xa9fe;
}
std::string hostName()
{
    wchar_t name[256]{}; DWORD size=256;
    return GetComputerNameW(name,&size) ? utf8(std::wstring(name,size)) : "KSwitcher";
}
std::vector<ULONG> broadcasts()
{
    ULONG size=16384; Bytes data(size);
    auto flags=GAA_FLAG_SKIP_ANYCAST|GAA_FLAG_SKIP_MULTICAST|GAA_FLAG_SKIP_DNS_SERVER;
    auto result=GetAdaptersAddresses(AF_INET,flags,nullptr,(IP_ADAPTER_ADDRESSES*)data.data(),&size);
    if(result==ERROR_BUFFER_OVERFLOW) { data.resize(size); result=GetAdaptersAddresses(AF_INET,flags,nullptr,(IP_ADAPTER_ADDRESSES*)data.data(),&size); }
    std::vector<ULONG> out;
    if(result!=NO_ERROR) return out;
    for(auto a=(IP_ADAPTER_ADDRESSES*)data.data();a;a=a->Next)
        if(a->OperStatus==IfOperStatusUp)
            for(auto u=a->FirstUnicastAddress;u;u=u->Next)
                if(u->Address.lpSockaddr->sa_family==AF_INET && u->OnLinkPrefixLength>0 && u->OnLinkPrefixLength<31)
                {
                    auto ip=((sockaddr_in*)u->Address.lpSockaddr)->sin_addr.s_addr;
                    if(localAddress(ip)) out.push_back(htonl(ntohl(ip) | (0xffffffffu >> u->OnLinkPrefixLength)));
                }
    return out;
}
}
struct LanService::Impl
{
    struct Peer { std::string name,epoch; sockaddr_in address{}; uint64_t seen{},verified{}; };
    std::filesystem::path path;
    std::string id,epoch=lan::randomHex(16),name=hostName(),status="Disabled.";
    std::map<std::string,Peer> peers;
    Json trusted=Json::object(),positions=Json::object();
    std::unique_ptr<lan::Pairing> pending;
    std::unique_ptr<lan::ClipboardExchange> clipboard;
    bool enabled{};
    std::atomic<bool> stopping{false};
    std::mutex mutex;
    std::thread worker;
    bool shareInput{};
    std::string inputPeer;
    int inputEdge=1;
    std::atomic<bool> inputConnected{false},abortInput{false};
    std::atomic<uint64_t> inputEpoch{0};
    std::atomic<uint64_t> inputConfiguration{0};
    std::atomic<bool> layoutChanged{false};
    uint64_t observedEpoch{},observedRevision{},observedConfiguration{};
    std::mutex inputMutex;
    std::vector<lan::RoutedInput> inputOut,inputIn;
    std::map<std::string,uint64_t> inputReady;
    std::unique_ptr<lan::InputMesh> router;
    std::string inputStatus="Input sharing disabled.";
    HWND inputWindow{};std::atomic<bool> inputNotification{false};
    explicit Impl(const std::filesystem::path &root,HWND window) : path(root/L"lan-pairing.bin"),inputWindow(window)
    {
        if(std::filesystem::exists(path))
        {
            auto j=Json::parse(unprotect(readFile(path,256*1024)));
            if(j.value("version",0)!=1 || !lan::identifier(j.value("id",""))) throw std::runtime_error("Invalid local LAN configuration.");
            id=j.at("id"); enabled=j.value("enabled",false); trusted=j.at("trusted");
            positions=j.value("positions",Json::object());
            shareInput=j.value("shareInput",false);inputPeer=j.value("inputPeer","");inputEdge=j.value("inputEdge",1);
            if(inputEdge<0||inputEdge>3||(shareInput&&trusted.empty()))throw std::runtime_error("Invalid input sharing configuration.");
            if(!trusted.is_object()||trusted.size()>64) throw std::runtime_error("Invalid paired computer list.");
            for(auto &[key,entry]:trusted.items())
                if(!lan::identifier(key)||key==id||entry.at("key").get<std::string>().size()!=64||
                   !entry.at("name").is_string()||entry.at("name").get<std::string>().size()>128)
                    throw std::runtime_error("Invalid paired computer key.");
        }
        else id=lan::randomHex(16);
        lan::fillPositions(positions,trusted,inputPeer,inputEdge);
        clipboard=std::make_unique<lan::ClipboardExchange>(id);
        configureRouter();
        if(enabled) start();
    }
    ~Impl() { stop(); }
    void save() { atomicWrite(path,protect(encoded({{"version",1},{"id",id},{"enabled",enabled},{"trusted",trusted},
        {"shareInput",shareInput},{"inputPeer",inputPeer},{"inputEdge",inputEdge},{"positions",positions}}))); }
    bool queue(const std::string &peer,lan::InputFrame frame)
    {
        std::lock_guard lock(inputMutex);
        if(!inputReady.contains(peer))return false;
        if(inputOut.size()>=512){abortInput=true;return false;}
        if(!inputOut.empty()&&inputOut.back().peer==peer&&inputOut.back().epoch==inputReady.at(peer)){
            std::vector<lan::InputFrame> tail{inputOut.back().frame};lan::enqueueInput(tail,frame);
            if(tail.size()==1){inputOut.back().frame=tail.front();return true;}
        }
        inputOut.push_back({peer,frame,inputReady.at(peer)});return true;
    }
    void configureRouter()
    {
        if(router)router->cancel(GetTickCount64());
        router.reset();++inputConfiguration;
        if(enabled&&shareInput){
            std::map<std::string,POINT> layout;for(auto &[peer,cell]:positions.items())if(trusted.contains(peer))layout[peer]={cell[0].get<LONG>(),cell[1].get<LONG>()};
            router=std::make_unique<lan::InputMesh>(id,lan::windowsInputEnvironment(),std::move(layout),[this](const auto &peer,auto f){return queue(peer,f);});
        }
        inputStatus=enabled&&shareInput?"Connecting input channels...":"Input sharing disabled.";
        observedRevision=0;
    }
    bool takeChange()
    {
        auto revision=router?router->revision():0;auto configuration=inputConfiguration.load();
        bool changed=revision!=observedRevision||configuration!=observedConfiguration;
        observedRevision=revision;observedConfiguration=configuration;return changed;
    }
    bool tickInput(bool allowEnter)
    {
        inputNotification=false;
        if(layoutChanged.exchange(false)){std::lock_guard lock(mutex);configureRouter();}
        if(!router)return false;
        auto epoch=inputEpoch.load();if(epoch!=observedEpoch){router->cancel(GetTickCount64());observedEpoch=epoch;}
        std::vector<lan::RoutedInput> incoming;std::map<std::string,uint64_t> ready;
        {std::lock_guard lock(inputMutex);incoming.swap(inputIn);ready=inputReady;}
        router->tick(ready,GetTickCount64());
        for(auto &frame:incoming)if(ready.contains(frame.peer)&&ready.at(frame.peer)==frame.epoch)
            router->receive(frame.peer,frame.frame,GetTickCount64(),allowEnter);
        return takeChange();
    }
    void start()
    {
        stopping=false; epoch=lan::randomHex(16); status="Searching for KSwitcher computers...";
        worker=std::thread([this]{ run(); });
    }
    void stop()
    {
        if(router)router->cancel(GetTickCount64());
        stopping=true;
        if(worker.joinable()) worker.join();
        inputConnected=false;++inputEpoch;
        {std::lock_guard lock(inputMutex);inputOut.clear();inputIn.clear();inputReady.clear();}
        std::lock_guard lock(mutex); pending.reset(); peers.clear();clipboard=std::make_unique<lan::ClipboardExchange>(id);
    }
    void send(SOCKET socket,const sockaddr_in &address,const Json &j)
    {
        auto bytes=j.dump(); if(bytes.size()>2048) return;
        sendto(socket,bytes.data(),int(bytes.size()),0,(const sockaddr*)&address,sizeof(address));
    }
    std::string heartbeat(const std::string &from,const std::string &to,const std::string &a,const std::string &b)
    { return "KSwitcher LAN v1|presence|"+from+"|"+to+"|"+a+"|"+b; }
    void receive(SOCKET socket,const sockaddr_in &address,const Json &j,uint64_t now)
    {
        if(j.value("app","")!="KSwitcher-LAN" || j.value("version",0)!=1) return;
        auto other=j.value("id",""); if(!lan::identifier(other)||other==id) return;
        auto type=j.value("type","");
        if(type=="hello")
        {
            auto display=j.value("name",""); auto session=j.value("epoch","");
            if(display.empty()||display.size()>128||!lan::identifier(session)||
               std::any_of(display.begin(),display.end(),[](unsigned char c){return c<32||c==127;})) return;
            // Discovery is untrusted and never grants input access.
            if(!peers.contains(other) && peers.size()>=64) return;
            auto &peer=peers[other];
            // A fresh announcement rotates the challenge every three seconds; it
            // does not revoke a recent successful proof for the same address.
            // Only a valid presence response renews verified, so untrusted/replayed
            // announcements cannot extend its PeerLifetime expiry.
            if(peer.address.sin_addr.s_addr!=address.sin_addr.s_addr) peer.verified=0;
            clipboard->peerSession(other,j.value("clipboardSession",""));
            peer.name=display; peer.epoch=session; peer.address=address; peer.seen=now;
            if(trusted.contains(other))
            {
                auto response=lan::envelope("presence",id); response["to"]=other; response["epoch"]=epoch;
                response["proof"]=lan::mac(trusted[other].at("key"),heartbeat(id,other,epoch,session));
                send(socket,address,response);
            }
            return;
        }
        if(!peers.contains(other)||j.value("to","")!=id) return;
        auto &peer=peers.at(other);
        if(peer.address.sin_addr.s_addr!=address.sin_addr.s_addr) return;
        if(type=="clipboard" && trusted.contains(other))
        {
            try {clipboard->receive(other,trusted.at(other).at("key"),j,now,GetClipboardSequenceNumber(),
                [&](const auto &target,const auto &packet){send(socket,peers.at(target).address,packet);});}
            catch(const std::exception &) {return;}
            if(!clipboard->deliveries.empty()&&inputWindow)PostMessageW(inputWindow,LanClipboardReadyMessage,0,0);
            return;
        }
        if(type=="presence" && trusted.contains(other) && j.value("epoch","")==peer.epoch &&
           lan::equal(j.value("proof",""),lan::mac(trusted[other].at("key"),heartbeat(other,id,peer.epoch,epoch))))
        { peer.verified=now; return; }
        if(pending && pending->peer==other)
        {
            for(const auto &response:pending->receive(j,now)) send(socket,address,response);
            if(pending->confirmed && (!trusted.contains(other)||trusted[other].at("key")!=pending->derived))
            {
                auto previous=trusted;auto previousPositions=positions;
                trusted[other]={{"name",peer.name},{"key",pending->derived}};
                try {lan::fillPositions(positions,trusted,inputPeer,inputEdge);save();} catch(...) {trusted=std::move(previous);positions=std::move(previousPositions);throw;}
                layoutChanged=true;peer.verified=now; status="Paired. Select a screen position and enable mouse and keyboard on both PCs.";
            }
        }
    }
    void run() noexcept
    {
        WSADATA startup{}; SOCKET socket=INVALID_SOCKET; bool winsock=false;
        try
        {
            if(WSAStartup(MAKEWORD(2,2),&startup)) throw std::runtime_error("Cannot start Windows networking.");
            winsock=true; socket=::socket(AF_INET,SOCK_DGRAM,IPPROTO_UDP);
            BOOL yes=TRUE; u_long nonblock=1;
            if(socket==INVALID_SOCKET || setsockopt(socket,SOL_SOCKET,SO_EXCLUSIVEADDRUSE,(char*)&yes,sizeof(yes)) ||
               setsockopt(socket,SOL_SOCKET,SO_BROADCAST,(char*)&yes,sizeof(yes)) || ioctlsocket(socket,FIONBIO,&nonblock))
                throw std::runtime_error("Cannot open LAN discovery socket.");
            int clipboardBuffer=1024*1024;
            setsockopt(socket,SOL_SOCKET,SO_RCVBUF,reinterpret_cast<const char*>(&clipboardBuffer),sizeof(clipboardBuffer));
            setsockopt(socket,SOL_SOCKET,SO_SNDBUF,reinterpret_cast<const char*>(&clipboardBuffer),sizeof(clipboardBuffer));
            sockaddr_in bound{}; bound.sin_family=AF_INET; bound.sin_port=htons(lan::Port);
            if(bind(socket,(sockaddr*)&bound,sizeof(bound))) throw std::runtime_error("UDP port 45883 is in use. Only one LAN-enabled KSwitcher may run per computer.");
            uint64_t nextHello=0,nextOffer=0;
            lan::InputHub hub(id);
            while(!stopping)
            {
                auto now=GetTickCount64();bool clipboardActive=false;
                std::vector<lan::InputEndpoint> endpoints;
                {
                    std::lock_guard lock(mutex);
                    clipboard->poll(trusted,now,[&](const auto &target,const auto &packet){
                        if(peers.contains(target))send(socket,peers.at(target).address,packet);
                    });
                    clipboardActive=clipboard->active();
                    if(shareInput)for(auto &[peer,entry]:trusted.items())
                        endpoints.push_back({peer,entry.at("key"),peers.contains(peer)?peers.at(peer).address.sin_addr.s_addr:0});
                    if(now>=nextHello)
                    {
                        epoch=lan::randomHex(16); // Fresh challenge prevents replayed presence from staying online.
                        auto hello=lan::envelope("hello",id); hello["name"]=name; hello["epoch"]=epoch;hello["clipboardSession"]=clipboard->session();
                        for(auto ip:broadcasts()) { auto address=bound; address.sin_addr.s_addr=ip; send(socket,address,hello); }
                        nextHello=now+3000;
                    }
                    for(auto it=peers.begin();it!=peers.end();) {
                        if(now-it->second.seen>lan::PeerLifetime) it=peers.erase(it); else ++it;
                    }
                    if(pending && pending->expired(now)) {
                        status=pending->confirmed ? "Pair saved. Select a computer to manage it." : "Pairing expired. Enter the same code and click Pair on both computers.";
                        pending.reset();
                    }
                    if(pending && now>=nextOffer && peers.contains(pending->peer)) {
                        send(socket,peers.at(pending->peer).address,pending->offer()); nextOffer=now+1000;
                    }
                }
                std::vector<lan::RoutedInput> out,in;
                {std::lock_guard lock(inputMutex);out.swap(inputOut);}
                if(abortInput.exchange(false)){hub.disconnect();out.clear();}
                hub.poll(endpoints,out,in,now);
                auto readyPeers=hub.ready();inputConnected=!readyPeers.empty();
                {std::lock_guard lock(inputMutex);inputReady=std::move(readyPeers);
                    for(auto &frame:in){
                        if(!inputIn.empty()&&inputIn.back().peer==frame.peer&&inputIn.back().epoch==frame.epoch){
                            std::vector<lan::InputFrame> tail{inputIn.back().frame};lan::enqueueInput(tail,frame.frame);
                            if(tail.size()==1){inputIn.back().frame=tail.front();continue;}
                        }
                        inputIn.push_back(frame);
                    }
                    if(inputIn.size()>512){inputIn.clear();abortInput=true;}
                }
                if(!in.empty()&&inputWindow&&!inputNotification.exchange(true))
                    if(!PostMessageW(inputWindow,LanInputReadyMessage,0,0))inputNotification=false;
                {std::lock_guard lock(mutex);inputStatus=hub.status();}
                fd_set read; FD_ZERO(&read); FD_SET(socket,&read); timeval delay{0,!endpoints.empty()||clipboardActive?5000:100000};
                int ready=select(0,&read,nullptr,nullptr,&delay);
                if(ready==SOCKET_ERROR) throw std::runtime_error("LAN discovery socket failed.");
                if(ready>0) for(int count=0;count<64;++count)
                {
                    char buffer[2049]; sockaddr_in sender{}; int size=sizeof(sender);
                    int received=recvfrom(socket,buffer,sizeof(buffer),0,(sockaddr*)&sender,&size);
                    if(received<=0) break;
                    if(received>2048||sender.sin_port!=htons(lan::Port)||!localAddress(sender.sin_addr.s_addr)) continue;
                    try { auto j=Json::parse(buffer,buffer+received); std::lock_guard lock(mutex); receive(socket,sender,j,GetTickCount64()); }
                    catch(const Json::exception &) {} // Malformed network data cannot interrupt discovery.
                }
            }
        }
        catch(const std::exception &e) { std::lock_guard lock(mutex); status=std::string("LAN stopped: ")+e.what(); }
        inputConnected=false;++inputEpoch;
        {std::lock_guard lock(mutex);if(shareInput)inputStatus="Input connection closed.";}
        if(socket!=INVALID_SOCKET) closesocket(socket);
        if(winsock) WSACleanup();
    }
    Bytes command(const Bytes &input)
    {
        auto j=Json::parse(input); auto action=j.value("action","status");
        if(action=="enable")
        {
            stop();
            bool previous=enabled; enabled=j.at("enabled").get<bool>();
            try {save();} catch(...) {enabled=previous;if(enabled) start();throw;}
            configureRouter();
            if(enabled) start(); else status="Disabled. No discovery sockets or worker thread.";
        }
        std::lock_guard lock(mutex);
        if(action=="position")
        {
            auto peer=j.at("peer").get<std::string>();int x=j.at("x"),y=j.at("y");
            if(!trusted.contains(peer))throw std::runtime_error("Pair this computer before arranging it.");
            lan::validateCell(positions,peer,x,y);
            auto before=positions;int beforeEdge=inputEdge;bool beforeActive=shareInput;
            positions[peer]=Json::array({x,y});
            if(peer==inputPeer){int edge=lan::cellEdge(positions[peer]);if(edge>=0)inputEdge=edge;}
            try{save();}catch(...){positions=std::move(before);inputEdge=beforeEdge;shareInput=beforeActive;throw;}
            configureRouter();
        }
        if(action=="input")
        {
            bool active=j.at("enabled");auto peer=j.value("peer",inputPeer);int edge=j.value("edge",inputEdge);
            if(edge<0||edge>3||(active&&(!enabled||trusted.empty()||(j.contains("peer")&&!trusted.contains(peer)))))throw std::runtime_error("Enable discovery and select a paired KSwitcher first.");
            auto beforePeer=inputPeer;auto beforePositions=positions;int beforeEdge=inputEdge;bool beforeActive=shareInput;
            if(j.contains("edge")&&trusted.contains(peer)){auto cell=lan::edgeCell(edge);lan::validateCell(positions,peer,cell[0],cell[1]);positions[peer]=cell;}
            shareInput=active;inputPeer=peer;inputEdge=edge;
            try{save();}catch(...){shareInput=beforeActive;inputPeer=beforePeer;inputEdge=beforeEdge;positions=std::move(beforePositions);throw;}
            configureRouter();inputStatus=active?"Connecting input channel...":"Input sharing disabled.";
        }
        if(action=="code") return encoded({{"code",lan::randomHex(16)}});
        if(action=="cancel") { pending.reset(); status="Pairing cancelled."; }
        if(action=="pair")
        {
            auto peer=j.at("peer").get<std::string>();
            if(!enabled||!peers.contains(peer)) throw std::runtime_error("Select an available KSwitcher computer first.");
            if(!trusted.contains(peer)&&trusted.size()>=64) throw std::runtime_error("Remove an old pair before adding another computer.");
            pending=std::make_unique<lan::Pairing>(id,peer,j.at("code").get<std::string>(),GetTickCount64());
            status="Waiting for the same code and Pair on the other KSwitcher (2 minutes)...";
        }
        if(action=="forget")
        {
            auto peer=j.at("peer").get<std::string>();auto oldTrusted=trusted;auto oldPositions=positions;bool oldSharing=shareInput;
            trusted.erase(peer);positions.erase(peer);if(trusted.empty())shareInput=false;
            try{save();}catch(...){trusted=std::move(oldTrusted);positions=std::move(oldPositions);shareInput=oldSharing;throw;}
            if(pending&&pending->peer==peer)pending.reset();
            configureRouter();
            status="Pair removed on this computer.";
        }
        Json result={{"enabled",enabled},{"id",id},{"name",name},{"status",status},{"pending",pending&&!pending->confirmed},{"peers",Json::array()},
            {"positions",positions},{"shareInput",shareInput},{"inputPeer",inputPeer},{"controlPeer",router?router->peer():std::string{}},{"inputEdge",inputEdge},{"inputConnected",inputConnected.load()},
            {"inputStatus",enabled&&shareInput?inputStatus:"Input sharing disabled."},{"controlMode",router?int(router->mode()):0}};
        {std::lock_guard inputLock(inputMutex);result["inputReady"]=inputReady;}
        auto now=GetTickCount64();
        for(const auto &[peer,p]:peers)
        {
            char ip[INET_ADDRSTRLEN]{}; InetNtopA(AF_INET,(void*)&p.address.sin_addr,ip,sizeof(ip));
            result["peers"].push_back({{"id",peer},{"name",p.name},{"address",ip},{"paired",trusted.contains(peer)},
                {"verified",p.verified && now-p.verified<lan::PeerLifetime},{"online",true}});
        }
        for(const auto &[peer,p]:trusted.items()) if(!peers.contains(peer))
            result["peers"].push_back({{"id",peer},{"name",p.at("name")},{"address",""},{"paired",true},{"verified",false},{"online",false}});
        return encoded(result);
    }
};
LanService::LanService(const std::filesystem::path &path,HWND window) : impl(std::make_unique<Impl>(path,window)) {}
LanService::~LanService()=default;
Bytes LanService::command(const Bytes &command) { return impl->command(command); }
bool LanService::clipboardEnabled()const
{
    std::lock_guard lock(impl->mutex);
    if(!impl->enabled)return false;
    for(auto &[id,entry]:impl->trusted.items())if(impl->peers.contains(id))return true;
    return false;
}
void LanService::publishClipboard(Bytes bytes,const std::string &app)
{
    std::lock_guard lock(impl->mutex);if(!impl->enabled)return;
    std::vector<std::string> targets;
    for(auto &[id,entry]:impl->trusted.items())if(impl->peers.contains(id))targets.push_back(id);
    if(!targets.empty())impl->clipboard->queue(targets,std::move(bytes),app,GetTickCount64());
}
std::vector<LanClipboard> LanService::takeClipboard()
{
    std::lock_guard lock(impl->mutex);std::vector<LanClipboard> result;
    for(auto &d:impl->clipboard->deliveries)if(impl->trusted.contains(d.peer))
        result.push_back({d.peer,impl->trusted.at(d.peer).at("name"),d.app,std::move(d.bytes),d.sequence});
    impl->clipboard->deliveries.clear();return result;
}
bool LanService::inputEnabled()const{return impl->enabled&&impl->shareInput;}
bool LanService::tickInput(bool allowEnter){return impl->tickInput(allowEnter);}
bool LanService::keyboard(const KBDLLHOOKSTRUCT &k,bool down){return impl->router&&impl->router->keyboard(k,down,GetTickCount64());}
bool LanService::mouse(UINT message,const MSLLHOOKSTRUCT &m,bool allowEnter){return impl->router&&impl->router->mouse(message,m,GetTickCount64(),allowEnter);}
bool LanService::takeInputChange(){return impl->takeChange();}
void LanService::cancelInput(){if(impl->router)impl->router->cancel(GetTickCount64());impl->abortInput=true;}
}