#include <winsock2.h>
#include <sddl.h>
#include "lan_files.h"
#include <bit>
#include <fstream>
#include <iostream>
#include <thread>

using namespace ks;
using namespace ks::lan;
namespace fs=std::filesystem;

namespace
{
void require(bool condition,const char *message)
{
    if(!condition) throw std::runtime_error(message);
}

void save(const fs::path &path,const std::string &text)
{
    std::ofstream file(path,std::ios::binary|std::ios::trunc);
    file.write(text.data(),text.size());
    require(bool(file),"write test fixture");
}

std::string load(const fs::path &path)
{
    std::ifstream file(path,std::ios::binary);
    return {std::istreambuf_iterator<char>(file),{}};
}

void verifyDirection(const FilesPeer *sourcePeer,const FilesPeer *targetPeer,
    const std::string &id,const std::string &key,const fs::path &root)
{
    auto connect=[&](const FilesPeer *peer)
    {
        return peer?std::make_unique<FilesConnection>(*peer,id,key):std::make_unique<FilesConnection>();
    };
    const auto sourcePath=root/L"source.bin";
    const auto targetPath=root/L"target.bin";
    std::string content(FileChunk*3+5,'a');
    for(size_t i=0;i<content.size();++i) content[i]=static_cast<char>((i/257)%251);
    save(sourcePath,content);
    save(targetPath,"original");
    auto source=connect(sourcePeer);
    const auto info=source->call(fileRequest("stat",sourcePath));
    checkFileReply(info);
    const auto size=info.at("size").get<uint64_t>(),modified=info.at("modified").get<uint64_t>();
    auto target=connect(targetPeer);
    try
    {
        streamFile(*source,*target,id,sourcePath,targetPath,size,modified,true,false,
            [](uint64_t done,uint64_t)
            {
                if(done==FileChunk) throw std::system_error(WSAECONNRESET,std::system_category());
                return true;
            });
        require(false,"interruption expected");
    }
    catch(const std::system_error &e){require(e.code().value()==WSAECONNRESET,"unexpected transfer failure");}
    target.reset();
    source.reset();
    require(load(targetPath)=="original","interruption overwrote target");
    require(load(sourcePath)==content,"interruption changed source");
    fs::path partial;
    for(const auto &entry:fs::directory_iterator(root))
        if(entry.path().extension()==L".partial") partial=entry.path();
    require(!partial.empty() && fs::file_size(partial)==FileChunk,"confirmed chunk was lost");

    source=connect(sourcePeer);
    target=connect(targetPeer);
    uint64_t offset=UINT64_MAX;
    streamFile(*source,*target,id,sourcePath,targetPath,size,modified,true,false,
        [&](uint64_t done,uint64_t){offset=std::min(offset,done);return true;});
    require(offset==FileChunk,"resumed at the wrong offset");
    require(load(targetPath)==content,"resumed file differs");
    require(!fs::exists(partial),"committed partial remains");
    require(load(sourcePath)==content,"copy removed source");

    save(targetPath,"original");
    try
    {
        streamFile(*source,*target,id,sourcePath,targetPath,size,modified,true,false,
            [](uint64_t done,uint64_t){return done<FileChunk;});
        require(false,"cancellation expected");
    }
    catch(const std::system_error &e){require(e.code().value()==ERROR_CANCELLED,"unexpected cancellation error");}
    require(load(targetPath)=="original" && !fs::exists(partial),"cancel touched original or retained partial");

    save(partial,std::string(FileChunk,'z'));
    uint64_t minimum=UINT64_MAX;
    streamFile(*source,*target,id,sourcePath,targetPath,size,modified,true,false,
        [&](uint64_t done,uint64_t){minimum=std::min(minimum,done);return true;});
    require(minimum==0 && load(targetPath)==content,"corrupt partial was reused");

    save(targetPath,"original");
    streamFile(*source,*target,id,sourcePath,targetPath,size,modified,true,true,{});
    require(load(targetPath)=="original"+content,"append did not preserve original prefix");
    checkFileReply(source->call(fileRequest("remove",sourcePath)));
    require(!fs::exists(sourcePath),"move source removal failed");
}
}

int main()
{
    WSADATA network{};
    if(WSAStartup(MAKEWORD(2,2),&network)) return 1;
    const auto root=fs::temp_directory_path()/wide("kswitcher-lan-files-"+randomHex(16));
    int result=0;
    try
    {
        fs::create_directory(root);
        const auto client=randomHex(16),first=randomHex(16),second=randomHex(16),key=randomHex(32);
        auto lookup=[&](const std::string &id){return id==client?key:std::string{};};
        FilesListener a(first,lookup,0),b(second,lookup,0);
        require(a.port() && b.port() && a.port()!=b.port(),"test listeners unavailable");
        std::jthread polling([&](std::stop_token stop)
        {
            while(!stop.stop_requested()){a.poll();b.poll();Sleep(5);}
        });
        const FilesPeer one{first,"127.0.0.1",L"One",a.port()},two{second,"127.0.0.1",L"Two",b.port()};
        for(const auto &direction:{std::pair<const FilesPeer*,const FilesPeer*>{nullptr,&one},
            {&one,nullptr},{&one,&two}})
            verifyDirection(direction.first,direction.second,client,key,root);

        try
        {
            FilesConnection unpaired(one,randomHex(16),key);
            require(false,"unpaired client authenticated");
        }
        catch(const std::system_error &){}
        FilesConnection host(one,client,key);
        const auto drives=host.call({{"action","drives"}});
        checkFileReply(drives);
        require(drives.at("drives").size()==static_cast<size_t>(std::popcount(GetLogicalDrives())),"drive mask differs");
        const auto denied=host.call(fileRequest("stat",root/L"NUL"));
        require(denied.at("error")==ERROR_INVALID_NAME,"reserved device accepted");
        // A deny ACE on file data applies to the resident's Windows token.
        const auto protectedFile=root/L"protected.bin";
        save(protectedFile,"restricted");
        const auto protectedInfo=host.call(fileRequest("stat",protectedFile));
        checkFileReply(protectedInfo);
        PSECURITY_DESCRIPTOR descriptor{};
        require(ConvertStringSecurityDescriptorToSecurityDescriptorW(
            L"D:(D;;0x1;;;WD)(A;;GA;;;WD)",SDDL_REVISION_1,&descriptor,nullptr)!=FALSE,
            "create test DACL");
        const auto applied=SetFileSecurityW(protectedFile.c_str(),DACL_SECURITY_INFORMATION,descriptor);
        LocalFree(descriptor);
        require(applied!=FALSE,"apply test DACL");
        auto protectedRead=fileRequest("openRead",protectedFile);
        protectedRead["size"]=protectedInfo.at("size");
        protectedRead["modified"]=protectedInfo.at("modified");
        require(host.call(protectedRead).at("error")==ERROR_ACCESS_DENIED,"Windows DACL did not deny remote read");
        checkFileReply(host.call(fileRequest("remove",protectedFile)));
        const auto tree=root/L"tree",file=tree/L"original.bin",copy=tree/L"copy.bin",renamed=tree/L"renamed.bin";
        checkFileReply(host.call(fileRequest("mkdir",tree)));
        save(file,"same host");
        auto operation=fileRequest("copy",file);
        operation["target"]=utf8(copy.native());
        operation["replace"]=false;
        operation["append"]=false;
        save(file,std::string(FileChunk*2,'s'));
        const auto largeInfo=host.call(fileRequest("stat",file));
        operation["size"]=largeInfo.at("size");
        operation["modified"]=largeInfo.at("modified");
        {
            FilesConnection interrupted(one,client,key);
            try
            {
                interrupted.call(operation,[](uint64_t done,uint64_t)
                {
                    if(done==FileChunk) throw std::system_error(WSAECONNRESET,std::system_category());
                    return true;
                });
                require(false,"same-host interruption expected");
            }
            catch(const std::system_error &e){require(e.code().value()==WSAECONNRESET,"same-host interruption error");}
        }
        require(!fs::exists(copy),"same-host interruption published a partial file");
        uint64_t firstOffset=UINT64_MAX;
        for(int i=0;i<100;++i)
        {
            const auto reply=host.call(operation,[&](uint64_t done,uint64_t)
            {
                firstOffset=std::min(firstOffset,done);
                return true;
            });
            if(reply.at("error")==ERROR_SHARING_VIOLATION){Sleep(10);continue;}
            checkFileReply(reply);
            break;
        }
        require(firstOffset==FileChunk && load(copy)==load(file),"same-host copy did not resume");
        operation=fileRequest("move",copy);
        operation["target"]=utf8(renamed.native());
        operation["replace"]=false;
        checkFileReply(host.call(operation));
        require(!fs::exists(copy) && load(renamed)==load(file),"same-host rename failed");
        checkFileReply(host.call(fileRequest("remove",renamed)));
        checkFileReply(host.call(fileRequest("remove",file)));
        checkFileReply(host.call(fileRequest("rmdir",tree)));
        require(!fs::exists(tree),"same-host directory remains");
    }
    catch(const std::exception &e){std::cerr<<e.what()<<'\n';result=1;}
    std::error_code ignored;
    fs::remove_all(root,ignored);
    WSACleanup();
    if(!result) std::cout<<"LAN file transfer regressions passed.\n";
    return result;
}