#include "lan.h"
#include "lan_pairing.h"
#include <iostream>
#include <tlhelp32.h>

namespace
{
int passed{};
void check(bool ok,const char *name) {if(!ok) throw std::runtime_error(name);++passed;}
template<class F> void rejects(F f,const char *name) {bool failed=false;try{f();}catch(...){failed=true;}check(failed,name);}
unsigned threads()
{
    HANDLE snapshot=CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD,0);
    if(snapshot==INVALID_HANDLE_VALUE) throw std::runtime_error("Cannot inspect test threads");
    THREADENTRY32 entry{sizeof(entry)};unsigned count=0;
    if(Thread32First(snapshot,&entry)) do {if(entry.th32OwnerProcessID==GetCurrentProcessId()) ++count;} while(Thread32Next(snapshot,&entry));
    CloseHandle(snapshot);return count;
}
}
int main(int argc,char **argv)
{
    using namespace ks;using namespace ks::lan;
    try
    {
        const std::string a(32,'a'),b(32,'b'),c(32,'c'),code="0123456789abcdef0123456789abcdef";
        check(normalizeCode("0123-4567-89AB-CDEF-0123-4567-89ab-cdef")==code,"code accepts groups and case");
        rejects([]{normalizeCode("123456");},"short guessable codes rejected");
        rejects([]{normalizeCode(std::string(32,'x'));},"invalid code rejected");
        check(identifier(randomHex(16)),"cryptographic instance ID");
        check(mac("key","The quick brown fox jumps over the lazy dog")=="f7bc83f430538424b13298e6aa6fb143ef4d59a14946175997479dbc2d1a3cd8", "HMAC SHA256 known vector");
        check(equal("abc","abc")&&!equal("abc","abd")&&!equal("abc","ab"),"proof comparison");
        Pairing left(a,b,code,100),right(b,a,code,100),wrong(b,a,std::string(32,'f'),100);
        check(wrong.receive(left.offer(),101).empty(),"different code cannot acknowledge");
        auto forged=left.offer();forged["id"]=c;
        check(right.receive(forged,101).empty(),"identity substitution rejected");
        forged=left.offer();forged["version"]=2;
        check(right.receive(forged,101).empty(),"unsupported protocol rejected");
        forged=left.offer();forged["to"]=c;
        check(right.receive(forged,101).empty(),"wrong destination rejected");
        auto ackR=right.receive(left.offer(),101),ackL=left.receive(right.offer(),101);
        check(ackR.size()==1&&ackL.size()==1,"mutual offers acknowledged");
        check(!left.confirmed&&!right.confirmed,"discovery and offers are not trust");
        auto confirmL=left.receive(ackR[0],102),confirmR=right.receive(ackL[0],102);
        check(confirmL.size()==1&&confirmR.size()==1,"authenticated challenges confirmed");
        left.receive(confirmR[0],103);right.receive(confirmL[0],103);
        check(left.confirmed&&right.confirmed&&left.derived==right.derived,"both endpoints derive same pair key");
        check(left.derived!=code&&left.derived.size()==64,"pair key separate from setup code");
        Pairing retry(a,b,code,200);
        check(retry.receive(ackR[0],201).empty()&&!retry.confirmed,"old ACK cannot pair new attempt");
        check(retry.receive(confirmR[0],201).empty()&&!retry.confirmed,"old confirmation cannot pair new attempt");
        Pairing expired(b,a,code,100);
        check(expired.receive(left.offer(),100+PairLifetime).empty(),"pairing window expires");
        check(left.receive(right.offer(),104).size()==1,"duplicate UDP offer is safely acknowledged");
        rejects([&]{Pairing self(a,a,code,0);},"self-pair rejected");
        auto root=std::filesystem::temp_directory_path()/wide("ks-lan-test-"+randomHex(16));
        std::filesystem::create_directory(root);
        struct Cleanup{std::filesystem::path p;~Cleanup(){std::error_code ec;std::filesystem::remove_all(p,ec);}} cleanup{root};
        auto request=[](LanService &service,Json j){auto s=j.dump();return Json::parse(service.command(Bytes(s.begin(),s.end())));};
        std::string identity;
        {
            auto idleThreads=threads();
            LanService service(root);
            check(threads()==idleThreads,"disabled LAN creates no worker thread");
            auto state=request(service,{{"action","status"}});
            check(!state.at("enabled").get<bool>()&&state.at("peers").empty(),"disabled by default");
            identity=state.at("id");
            check(!std::filesystem::exists(root/L"lan-pairing.bin"),"viewing disabled status does not create config");
            auto generated=request(service,{{"action","code"}}).at("code").get<std::string>();
            check(identifier(generated),"generated pairing code has 128 bits");
            rejects([&]{request(service,{{"action","pair"},{"peer",b},{"code",code}});},"cannot pair undiscovered peer while disabled");
            request(service,{{"action","enable"},{"enabled",false}});
            auto sealed=readFile(root/L"lan-pairing.bin",256*1024);
            check(std::string(sealed.begin(),sealed.end()).find(identity)==std::string::npos,"local configuration protected by DPAPI");
            if(argc>1&&std::string(argv[1])=="--network")
            {
                request(service,{{"action","enable"},{"enabled",true}});
                Sleep(400);
                auto active=request(service,{{"action","status"}});
                check(active.at("enabled").get<bool>() && active.at("status").get<std::string>().find("LAN stopped:")==std::string::npos,"opt-in discovery starts successfully");
                check(threads()>idleThreads,"enabled discovery has a worker");
                request(service,{{"action","enable"},{"enabled",false}});
                check(threads()==idleThreads,"disabling discovery joins its worker");
                // Reopening immediately exercises socket cleanup and exclusive bind.
                request(service,{{"action","enable"},{"enabled",true}});Sleep(200);
                check(request(service,{{"action","status"}}).at("status").get<std::string>().find("LAN stopped:")==std::string::npos,"discovery port released and reusable");
                request(service,{{"action","enable"},{"enabled",false}});
            }
        }
        {
            LanService restored(root);auto state=request(restored,{{"action","status"}});
            check(state.at("id")==identity&&!state.at("enabled").get<bool>(),"stable local instance ID and disabled state restored");
        }
        {
            auto saved=Json::parse(unprotect(readFile(root/L"lan-pairing.bin",256*1024)));
            saved["trusted"][b]={{"name","Other KSwitcher"},{"key",left.derived}};
            auto json=saved.dump();atomicWrite(root/L"lan-pairing.bin",protect(Bytes(json.begin(),json.end())));
            LanService restored(root);auto state=request(restored,{{"action","status"}});
            check(state["peers"].size()==1 && state["peers"][0].at("paired").get<bool>() && !state["peers"][0].at("online").get<bool>(),"saved pair remains visible offline");
            check(state.dump().find(left.derived)==std::string::npos,"status never exposes stored pair keys");
            rejects([&]{request(restored,{{"action","input"},{"enabled",true},{"peer",b},{"edge",1}});},"input cannot be enabled without discovery");
            rejects([&]{request(restored,{{"action","input"},{"enabled",false},{"edge",4}});},"invalid screen edge rejected");
            auto position=request(restored,{{"action","input"},{"enabled",false},{"peer",b},{"edge",2}});
            check(position.at("inputEdge")==2&&!restored.inputEnabled(),"screen position saved without enabling input");
            if(argc>1&&std::string(argv[1])=="--network")
            {
                request(restored,{{"action","enable"},{"enabled",true}});
                rejects([&]{request(restored,{{"action","input"},{"enabled",true},{"peer",c},{"edge",1}});},"discovered but unpaired peer cannot receive input");
                request(restored,{{"action","input"},{"enabled",true},{"peer",b},{"edge",2}});
                check(restored.inputEnabled(),"paired input enabled with explicit screen position");
                request(restored,{{"action","enable"},{"enabled",false}});
                check(!restored.inputEnabled(),"disabling discovery also stops input runtime");
            }
            request(restored,{{"action","forget"},{"peer",b}});
            check(request(restored,{{"action","status"}})["peers"].empty(),"forget removes offline pair");
            check(!request(restored,{{"action","status"}}).at("shareInput").get<bool>(),"forgetting the input peer revokes persisted input permission");
        }
        if(argc>1&&std::string(argv[1])=="--network")
        {
            auto saved=Json::parse(unprotect(readFile(root/L"lan-pairing.bin",256*1024)));
            saved["trusted"][b]={{"name","Restart peer"},{"key",left.derived}};
            saved["enabled"]=true;saved["shareInput"]=true;saved["inputPeer"]=b;saved["inputEdge"]=3;
            auto text=saved.dump();atomicWrite(root/L"lan-pairing.bin",protect(Bytes(text.begin(),text.end())));
            for(int restart=0;restart<2;++restart)
            {
                LanService restored(root);Sleep(100);
                auto state=request(restored,{{"action","status"}});
                check(restored.inputEnabled()&&state.at("enabled")==true&&state.at("inputPeer")==b&&state.at("inputEdge")==3,
                    "restart restores discovery, pairing, input permission and screen edge without opening settings");
            }
        }
        {
            auto folder=root/L"layout";std::filesystem::create_directories(folder);
            auto saved=Json::parse(unprotect(readFile(root/L"lan-pairing.bin",256*1024)));
            saved["enabled"]=false;saved["shareInput"]=false;saved["inputPeer"]=b;saved["inputEdge"]=1;saved.erase("positions");
            saved["trusted"]={{b,{{"name","Right PC"},{"key",left.derived}}},{c,{{"name","Third PC"},{"key",left.derived}}}};
            auto text=saved.dump();atomicWrite(folder/L"lan-pairing.bin",protect(Bytes(text.begin(),text.end())));
            {
                LanService layout(folder);auto state=request(layout,{{"action","status"}});
                check(state["positions"][b]==Json::array({1,0}),"legacy side migrates to monitor grid");
                check(state["positions"].size()==2&&state["positions"][c]!=state["positions"][b],"three-PC layout has unique cells");
                request(layout,{{"action","position"},{"peer",c},{"x",-1},{"y",0}});
                request(layout,{{"action","position"},{"peer",b},{"x",0},{"y",-1}});
                rejects([&]{request(layout,{{"action","position"},{"peer",c},{"x",0},{"y",-1}});},"occupied cell rejected");
                rejects([&]{request(layout,{{"action","position"},{"peer",c},{"x",0},{"y",0}});},"local monitor cell protected");
                rejects([&]{request(layout,{{"action","position"},{"peer",c},{"x",9},{"y",0}});},"out-of-grid position rejected");
                check(request(layout,{{"action","status"}})["positions"][c]==Json::array({-1,0}),"failed drops preserve previous position");
            }
            {
                LanService layout(folder);auto state=request(layout,{{"action","status"}});
                check(state["positions"][b]==Json::array({0,-1})&&state["positions"][c]==Json::array({-1,0}),"all monitor positions survive restart");
                check(state["inputEdge"]==2,"active peer direction follows monitor position");
                if(argc>1&&std::string(argv[1])=="--network"){
                    request(layout,{{"action","enable"},{"enabled",true}});
                    request(layout,{{"action","input"},{"enabled",true}});
                    request(layout,{{"action","position"},{"peer",c},{"x",0},{"y",-2}});
                    auto multi=request(layout,{{"action","status"}});
                    check(multi["shareInput"]==true&&multi["positions"][c]==Json::array({0,-2}),"multi-PC sharing remains enabled for a chained nonadjacent destination");
                    request(layout,{{"action","forget"},{"peer",b}});
                    check(request(layout,{{"action","status"}})["shareInput"]==true,"forgetting legacy selected pair keeps remaining pairs enabled");
                }
                request(layout,{{"action","forget"},{"peer",c}});
                check(!request(layout,{{"action","status"}})["positions"].contains(c),"forgetting pair removes its monitor");
            }
        }
        std::cout<<"Passed "<<passed<<" LAN discovery/pairing checks\n";
    }
    catch(const std::exception &e){std::cerr<<"FAIL: "<<e.what()<<'\n';return 1;}
}