#include <winsock2.h>
#include "lan_input.h"
#include <iostream>
#include <deque>

using namespace ks;using namespace ks::lan;
namespace
{
int passed{};
void check(bool value,const char *text){if(!value)throw std::runtime_error(text);++passed;}
template<class F>void rejects(F f,const char *text){bool failed=false;try{f();}catch(...){failed=true;}check(failed,text);}
struct Desktop
{
    RECT rect{-1920,-200,0,880};POINT point{-1000,300};bool available=true,physicalHeld=false,allow=true;
    std::vector<INPUT> events;
    InputEnvironment environment()
    {
        return {[this]{return rect;},[this]{return point;},[this](const INPUT &i){
            if(!allow)return false;
            events.push_back(i);
            if(i.type==INPUT_MOUSE&&(i.mi.dwFlags&MOUSEEVENTF_MOVE))
                point={rect.left+LONG(int64_t(i.mi.dx)*(rect.right-rect.left)/65536),rect.top+LONG(int64_t(i.mi.dy)*(rect.bottom-rect.top)/65536)};
            return true;
        },[this]{return available;},[this]{return physicalHeld;}};
    }
    unsigned keyEvents(bool down)const
    {return unsigned(std::count_if(events.begin(),events.end(),[down](auto &i){return i.type==INPUT_KEYBOARD&&bool(i.ki.dwFlags&KEYEVENTF_KEYUP)!=down;}));}
};
KBDLLHOOKSTRUCT key(unsigned vk,unsigned scan=0){KBDLLHOOKSTRUCT k{};k.vkCode=vk;k.scanCode=scan;return k;}
bool motion(InputRouter &router,Desktop &desktop,POINT p,uint64_t now)
{desktop.point=p;MSLLHOOKSTRUCT m{};m.pt=p;return router.mouse(WM_MOUSEMOVE,m,now);}
void exchange(InputRouter &source,InputRouter &target,std::deque<InputFrame> &out,uint64_t now)
{(void)source;while(!out.empty()){auto f=out.front();out.pop_front();target.receive(f,now);}}
struct Pair
{
    Desktop a,b;std::deque<InputFrame> ab,ba;
    InputRouter left,right;
    explicit Pair(int side):left(a.environment(),[this](auto f){ab.push_back(f);return true;},side,true),
        right(b.environment(),[this](auto f){ba.push_back(f);return true;},side^1,false)
    {b.rect={0,0,1600,900};b.point={600,400};left.tick(true,1000);right.tick(true,1000);}
    void deliver(uint64_t now=1100){exchange(left,right,ab,now);exchange(right,left,ba,now);}
    void enter(int edge)
    {
        motion(left,a,edgePoint(a.rect,edge,30000,0),1000);
        check(left.mode()==ControlMode::Requesting,"crossing edge requests ownership, not immediate suppression");
        deliver();check(left.mode()==ControlMode::Sending&&right.mode()==ControlMode::Receiving,"mouse and keyboard ownership acknowledged together");
    }
};
void routerChecks()
{
    for(int side=0;side<4;++side)
    {
        Pair p(side);auto middle=edgePoint(p.a.rect,side,32768,0);
        check(atEdge(middle,p.a.rect,side),"all four virtual-screen edges supported");
        check(std::abs(edgePosition(middle,p.a.rect,side)-32768)<100,"negative monitor coordinates normalize");
        p.enter(side);
        auto old=p.b.point;
        for(int step=0;step<20;++step)
        {
            auto point=p.a.point;if(side==0)--point.x;else if(side==1)++point.x;else if(side==2)--point.y;else ++point.y;
            motion(p.left,p.a,point,1110+step);p.deliver(1110+step);
        }
        check(side<2 ? std::abs(p.b.point.x-old.x)==20 : std::abs(p.b.point.y-old.y)==20,"one-pixel mouse movements have no accumulating warp bias");
        check(p.left.keyboard(key('A',0x1e),true,1200),"source key down suppressed");p.deliver(1200);
        check(p.a.keyEvents(true)==0&&p.b.keyEvents(true)==1,"key only injected on computer with cursor");
        check(p.b.events.back().ki.dwFlags&KEYEVENTF_SCANCODE,"physical scan code preserved across layouts");
        check(p.left.keyboard(key('A',0x1e),false,1210),"source key up suppressed");p.deliver(1210);
        check(p.b.keyEvents(false)==1,"remote key released");
        auto injected=key('A',0x1e);injected.flags=LLKHF_INJECTED;injected.dwExtraInfo=remoteInputMarker();
        check(!p.right.keyboard(injected,true,1220),"received input continues to destination hotkeys/snippets, no routing loop");
        MSLLHOOKSTRUCT click{};check(p.left.mouse(WM_LBUTTONDOWN,click,1230),"source click suppressed");p.deliver(1230);
        check(p.b.events.back().mi.dwFlags==MOUSEEVENTF_LEFTDOWN,"remote left button pressed");
        p.left.mouse(WM_LBUTTONUP,click,1240);p.deliver(1240);
        check(p.b.events.back().mi.dwFlags==MOUSEEVENTF_LEFTUP,"remote left button released");
        click.mouseData=DWORD(120)<<16;p.left.mouse(WM_MOUSEWHEEL,click,1250);p.deliver(1250);
        check(p.b.events.back().mi.dwFlags==MOUSEEVENTF_WHEEL&&p.b.events.back().mi.mouseData==120,"wheel delta forwarded");
        auto centre=p.a.point;POINT moved=centre;
        if(side==0)moved.x-=30;else if(side==1)moved.x+=30;else if(side==2)moved.y-=30;else moved.y+=30;
        check(motion(p.left,p.a,moved,1300),"remote mouse move suppressed locally");p.deliver(1300);
        check(p.b.point.x!=old.x||p.b.point.y!=old.y,"cursor moves on target");
        // Repeated source movement towards the return edge.
        for(int n=0;n<5&&p.left.mode()==ControlMode::Sending;++n)
        {
            moved=p.a.point;if(side==0)moved.x+=400;else if(side==1)moved.x-=400;else if(side==2)moved.y+=400;else moved.y-=400;
            motion(p.left,p.a,moved,1400+n);p.deliver(1400+n);
        }
        check(p.left.mode()==ControlMode::Local&&p.right.mode()==ControlMode::Local,"crossing target return edge restores local ownership");
        check(!atEdge(p.a.point,p.a.rect,side),"return cursor inset avoids immediate bounce");
        check(!p.left.keyboard(key('B'),true,1500),"normal local typing resumes");p.left.keyboard(key('B'),false,1501);
    }
    {
        Pair p(1);p.enter(1);p.left.keyboard(key(VK_LCONTROL),true,1200);p.left.keyboard(key(VK_LMENU),true,1201);p.deliver();
        check(p.left.keyboard(key(VK_ESCAPE),true,1202),"emergency return chord consumed");p.deliver();
        check(p.left.mode()==ControlMode::Local&&p.right.mode()==ControlMode::Local,"emergency chord returns both sides");
        check(p.b.keyEvents(false)==2,"emergency return releases remote modifiers");
        check(p.left.keyboard(key(VK_LCONTROL),false,1203)&&p.left.keyboard(key(VK_LMENU),false,1204),"stray source modifier releases consumed after return");
    }
    {
        Pair p(1);p.enter(1);p.left.keyboard(key('A'),true,1200);p.deliver();
        p.left.tick(false,1300);p.right.tick(false,1300);
        check(p.b.keyEvents(false)==1&&p.left.mode()==ControlMode::Local,"disconnect releases held remote key");
        check(p.left.keyboard(key('A'),false,1400),"source does not leak orphaned key-up after disconnect");
    }
    {
        Pair p(1);p.enter(1);p.left.keyboard(key('A'),true,1200);p.deliver();p.b.available=false;p.right.tick(true,1300);
        check(p.right.mode()==ControlMode::Local&&p.b.keyEvents(false)==0,"secure desktop stops ownership without injecting there");
        p.b.available=true;p.right.tick(true,1400);check(p.b.keyEvents(false)==1,"blocked release retried when normal desktop returns");
    }
    {
        Pair p(1);p.a.physicalHeld=true;motion(p.left,p.a,edgePoint(p.a.rect,1,30000,0),1000);
        check(p.left.mode()==ControlMode::Local,"no cross-screen handoff with physically held key/button");
        p.a.physicalHeld=false;motion(p.left,p.a,edgePoint(p.a.rect,1,30000,0),1100);
        p.left.keyboard(key('X'),true,1101);p.deliver(1200);
        check(p.left.mode()==ControlMode::Local&&p.right.mode()==ControlMode::Local,"typing during handoff cancels stale accept");
    }
    {
        Pair p(1);motion(p.left,p.a,edgePoint(p.a.rect,1,30000,0),1000);motion(p.right,p.b,edgePoint(p.b.rect,0,30000,0),1000);
        p.deliver();p.deliver();check(p.left.mode()==ControlMode::Sending&&p.right.mode()==ControlMode::Receiving,"simultaneous crossing has deterministic ownership");
    }
    {
        Pair p(1);p.enter(1);p.left.keyboard(key('A'),true,1200);p.deliver();
        check(!p.right.keyboard(key('B'),true,1300),"local physical keyboard takes over target");p.deliver();
        check(p.left.mode()==ControlMode::Local&&p.b.keyEvents(false)==1,"local takeover returns source and releases remote keys");
    }
    {
        Pair p(1);p.b.allow=false;motion(p.left,p.a,edgePoint(p.a.rect,1,30000,0),1000);p.deliver();
        check(p.left.mode()==ControlMode::Local&&p.right.mode()==ControlMode::Local,"blocked SendInput fails closed without capturing source");
    }
    {
        Pair p(1);p.enter(1);MSLLHOOKSTRUCT m{};m.mouseData=DWORD(XBUTTON2)<<16;
        p.left.mouse(WM_XBUTTONDOWN,m,1200);p.deliver();p.right.tick(false,1300);
        check(p.b.events.back().mi.dwFlags==MOUSEEVENTF_XUP&&p.b.events.back().mi.mouseData==XBUTTON2,"disconnect releases held extra mouse button");
    }
    {
        Pair p(1);motion(p.left,p.a,edgePoint(p.a.rect,1,30000,0),1000);p.left.tick(true,2101);
        check(p.left.mode()==ControlMode::Local,"unanswered ownership request times out locally");p.deliver(2200);p.deliver(2200);
        check(p.left.mode()==ControlMode::Local&&p.right.mode()==ControlMode::Local,"late handoff acknowledgement cannot steal control");
    }
    {
        Desktop d;bool accept=true;std::deque<InputFrame> frames;
        InputRouter router(d.environment(),[&](auto f){frames.push_back(f);return accept;},1,true);
        router.tick(true,1000);motion(router,d,edgePoint(d.rect,1,30000,0),1000);
        auto token=frames.front().token;router.receive({InputType::Accept,token},1100);
        accept=false;check(router.keyboard(key('Q'),true,1200),"failed queue still consumes the key rather than typing locally");
        check(router.mode()==ControlMode::Local&&router.keyboard(key('Q'),false,1300),"queue failure returns control and consumes orphan release");
    }
}
uint16_t freePort()
{
    SOCKET s=socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);if(s==INVALID_SOCKET)throw std::runtime_error("test socket");
    sockaddr_in a{};a.sin_family=AF_INET;a.sin_addr.s_addr=htonl(INADDR_LOOPBACK);
    if(bind(s,(sockaddr*)&a,sizeof(a))){closesocket(s);throw std::runtime_error("test bind");}
    int size=sizeof(a);getsockname(s,(sockaddr*)&a,&size);auto port=ntohs(a.sin_port);closesocket(s);return port;
}
void channelChecks()
{
    std::string a(32,'a'),b(32,'b'),secret(64,'f');auto na=randomHex(16),nb=randomHex(16);
    auto k=inputSessionKey(secret,a,b,na,nb);check(k!=inputSessionKey(secret,b,a,nb,na),"directional session keys differ");
    check(inputSessionKey(secret,a,b,na,nb,0)!=inputSessionKey(secret,a,b,na,nb,1),"screen geometry bound into authenticated session");
    InputCipher tx(k),rx(k);InputFrame f{InputType::Key,37,'A',0x1e,1};auto packet=tx.encrypt(f);
    check(packet.size()==36&&rx.decrypt(packet)==f,"AES-GCM authenticated input roundtrip");
    rejects([&]{rx.decrypt(packet);},"encrypted frame replay rejected");
    auto next=tx.encrypt({InputType::Move,37,-12,400});auto tampered=next;tampered[4]^=1;
    rejects([&]{rx.decrypt(tampered);},"ciphertext modification rejected");
    check(rx.decrypt(next)==InputFrame{InputType::Move,37,-12,400},"failed authentication does not expose or advance data");
    InputCipher wrong(inputSessionKey(secret,b,a,nb,na));rejects([&]{wrong.decrypt(packet);},"reflection across directions rejected");
    InputCipher otherSession(inputSessionKey(secret,a,b,randomHex(16),nb));rejects([&]{otherSession.decrypt(packet);},"old-session input rejected");
    WSADATA data{};if(WSAStartup(MAKEWORD(2,2),&data))throw std::runtime_error("WSAStartup");
    struct Close{~Close(){WSACleanup();}} close;
    {
        auto port=freePort();InputChannel client(a,b,secret,1,port),server(b,a,secret,0,port);
        auto deadline=GetTickCount64()+5000;std::vector<InputFrame> received;
        while((!client.ready()||!server.ready())&&GetTickCount64()<deadline)
        {
            server.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());client.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());Sleep(1);
        }
        check(client.ready()&&server.ready(),"real TCP peers mutually authenticate with paired key");
        check(received.empty(),"handshake produces no input events");
        client.poll(htonl(INADDR_LOOPBACK),{f},received,GetTickCount64());
        for(int n=0;n<100&&received.empty();++n){server.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());client.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());Sleep(1);}
        check(received.size()==1&&received[0]==f,"real TCP channel delivers one ordered authenticated key");received.clear();
        InputFrame reverse{InputType::Wheel,37,-120,1};server.poll(htonl(INADDR_LOOPBACK),{reverse},received,GetTickCount64());
        for(int n=0;n<100&&received.empty();++n){client.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());Sleep(1);}
        check(received.size()==1&&received[0]==reverse,"real input channel is bidirectional");
        auto generation=client.generation();client.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64()+3000);
        check(!client.ready()&&client.generation()!=generation,"heartbeat timeout closes channel and changes generation");
    }
    {
        auto port=freePort();InputChannel client(a,b,secret,1,port),server(b,a,std::string(64,'e'),0,port);
        std::vector<InputFrame> received;auto start=GetTickCount64();
        while(GetTickCount64()-start<400){server.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());client.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());Sleep(1);}
        check(!client.ready()&&!server.ready()&&received.empty(),"wrong saved pair key never authorizes input");
    }
    {
        auto port=freePort();InputChannel client(a,b,secret,1,port),server(b,a,secret,1,port);
        std::vector<InputFrame> received;auto start=GetTickCount64();
        while(GetTickCount64()-start<400){server.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());client.poll(htonl(INADDR_LOOPBACK),{},received,GetTickCount64());Sleep(1);}
        check(!client.ready()&&!server.ready()&&received.empty(),"conflicting screen arrangement rejected before input");
    }
}
}
int main()
{
    try{channelChecks();routerChecks();std::cout<<"Passed "<<passed<<" LAN input checks\n";}
    catch(const std::exception &e){std::cerr<<"FAIL: "<<e.what()<<'\n';return 1;}
}