#include "lan_input.h"
#include <algorithm>
#include <bit>
#include <atomic>

namespace ks::lan
{
ULONG_PTR remoteInputMarker(){return inputMarker() ^ ULONG_PTR(0x4c414e49);}
ULONG_PTR cursorWarpMarker(){return inputMarker() ^ ULONG_PTR(0x4c414e57);}
namespace
{
uint32_t ownershipToken()
{
    static std::atomic<uint32_t> next{1};
    auto token=next.fetch_add(1);return token?token:next.fetch_add(1);
}
HCURSOR parkedCursor()
{
    static HCURSOR cursor=[] {unsigned char mask[128],pixels[128]{};
        std::fill(std::begin(mask),std::end(mask),0xff);
        return CreateCursor(GetModuleHandleW(nullptr),0,0,32,32,mask,pixels);}();
    return cursor;
}
int monitorDpi(POINT point)
{
    static HMODULE module=LoadLibraryW(L"shcore.dll");
    using Query=HRESULT(WINAPI *)(HMONITOR,int,UINT*,UINT*);
    static Query query=module?std::bit_cast<Query>(GetProcAddress(module,"GetDpiForMonitor")):nullptr;
    UINT x=96,y=96;if(query)query(MonitorFromPoint(point,MONITOR_DEFAULTTONEAREST),0,&x,&y);
    return int(std::clamp<UINT>(x,48,768));
}
LRESULT CALLBACK controlOverlay(HWND w,UINT message,WPARAM wp,LPARAM lp)
{
    if(message==WM_NCCREATE)SetWindowLongPtrW(w,GWLP_USERDATA,(LONG_PTR)((CREATESTRUCTW*)lp)->lpCreateParams);
    if(message==WM_SETCURSOR){SetCursor(GetWindowLongPtrW(w,GWLP_USERDATA)?parkedCursor():LoadCursorW(nullptr,IDC_ARROW));return TRUE;}
    if(message==WM_MOUSEACTIVATE)return MA_NOACTIVATE;
    if(message==WM_PAINT)
    {
        PAINTSTRUCT paint{};HDC dc=BeginPaint(w,&paint);RECT r{};GetClientRect(w,&r);
        FillRect(dc,&r,(HBRUSH)(COLOR_INFOBK+1));SetBkMode(dc,TRANSPARENT);SetTextColor(dc,GetSysColor(COLOR_INFOTEXT));
        auto previous=SelectObject(dc,GetStockObject(DEFAULT_GUI_FONT));
        DrawTextW(dc,L"KSwitcher: mouse + keyboard on the other PC\nCtrl+Alt+Esc returns control here",-1,&r,DT_CENTER|DT_WORDBREAK);
        SelectObject(dc,previous);EndPaint(w,&paint);return 0;
    }
    return DefWindowProcW(w,message,wp,lp);
}
struct ControlOverlay
{
    HWND notice{},cover{};HCURSOR previous{};
    ~ControlOverlay(){hide();}
    void hide(){bool shown=notice||cover;if(cover)DestroyWindow(cover);if(notice)DestroyWindow(notice);notice=cover=nullptr;
        if(shown)SetCursor(previous?previous:LoadCursorW(nullptr,IDC_ARROW));
        previous=nullptr;}
    void show()
    {
        if(notice)return;
        WNDCLASSW cls{};cls.hInstance=GetModuleHandleW(nullptr);cls.lpfnWndProc=controlOverlay;cls.lpszClassName=L"KSwitcher.RemoteControlOverlay";cls.hCursor=LoadCursorW(nullptr,IDC_ARROW);
        RegisterClassW(&cls);
        auto style=WS_EX_TOPMOST|WS_EX_TOOLWINDOW|WS_EX_NOACTIVATE;
        POINT cursor{};GetCursorPos(&cursor);MONITORINFO monitor{sizeof(monitor)};GetMonitorInfoW(MonitorFromPoint(cursor,MONITOR_DEFAULTTONEAREST),&monitor);
        CURSORINFO ci{sizeof(ci)};if(GetCursorInfo(&ci))previous=ci.hCursor;
        int size=MulDiv(64,monitorDpi(cursor),96);
        cover=CreateWindowExW(style|WS_EX_LAYERED,cls.lpszClassName,L"",WS_POPUP,
            (monitor.rcMonitor.left+monitor.rcMonitor.right-size)/2,(monitor.rcMonitor.top+monitor.rcMonitor.bottom-size)/2,
            size,size,nullptr,nullptr,cls.hInstance,(void*)1);
        if(cover){SetLayeredWindowAttributes(cover,0,1,LWA_ALPHA);ShowWindow(cover,SW_SHOWNOACTIVATE);}
        notice=CreateWindowExW(style,cls.lpszClassName,L"KSwitcher remote control",WS_POPUP,monitor.rcMonitor.left+20,monitor.rcMonitor.top+20,400,48,nullptr,nullptr,cls.hInstance,nullptr);
        if(notice)ShowWindow(notice,SW_SHOWNOACTIVATE);
        if(cover)SetCursor(parkedCursor());
    }
};
bool ordinaryDesktop()
{
    HDESK d=OpenInputDesktop(0,FALSE,DESKTOP_READOBJECTS);if(!d)return false;
    wchar_t name[256]{};DWORD count{};bool ok=GetUserObjectInformationW(d,UOI_NAME,name,sizeof(name),&count)!=0;
    CloseDesktop(d);return ok&&_wcsicmp(name,L"Default")==0;
}
INPUT button(unsigned index,bool down)
{
    INPUT input{};input.type=INPUT_MOUSE;
    static constexpr DWORD flags[5][2]={{MOUSEEVENTF_LEFTUP,MOUSEEVENTF_LEFTDOWN},{MOUSEEVENTF_RIGHTUP,MOUSEEVENTF_RIGHTDOWN},
        {MOUSEEVENTF_MIDDLEUP,MOUSEEVENTF_MIDDLEDOWN},{MOUSEEVENTF_XUP,MOUSEEVENTF_XDOWN},{MOUSEEVENTF_XUP,MOUSEEVENTF_XDOWN}};
    input.mi.dwFlags=flags[index][down];if(index>=3)input.mi.mouseData=index==3?XBUTTON1:XBUTTON2;return input;
}
}
InputEnvironment windowsInputEnvironment()
{
    auto overlay=std::make_shared<ControlOverlay>();
    return {
        []{LONG x=GetSystemMetrics(SM_XVIRTUALSCREEN),y=GetSystemMetrics(SM_YVIRTUALSCREEN);return RECT{x,y,x+GetSystemMetrics(SM_CXVIRTUALSCREEN),y+GetSystemMetrics(SM_CYVIRTUALSCREEN)};},
        []{POINT p{};GetCursorPos(&p);return p;},
        [](const INPUT &i){INPUT input=i;return SendInput(1,&input,sizeof(input))==1;},
        ordinaryDesktop,
        []{for(int key=1;key<256;++key)if(GetAsyncKeyState(key)&0x8000)return true;return false;},
        []{POINT p{};GetCursorPos(&p);MONITORINFO info{sizeof(info)};
            if(GetMonitorInfoW(MonitorFromPoint(p,MONITOR_DEFAULTTONEAREST),&info))p={(info.rcMonitor.left+info.rcMonitor.right)/2,(info.rcMonitor.top+info.rcMonitor.bottom)/2};
            return p;},
        [overlay](ControlMode mode){if(mode==ControlMode::Sending)overlay->show();else overlay->hide();},
        monitorDpi
    };
}
bool atEdge(POINT p,RECT r,int side)
{return side==0?p.x<=r.left:side==1?p.x>=r.right-1:side==2?p.y<=r.top:p.y>=r.bottom-1;}
int edgePosition(POINT p,RECT r,int side)
{
    auto length=side<2?r.bottom-r.top:r.right-r.left;
    auto offset=side<2?p.y-r.top:p.x-r.left;
    return int(std::clamp<int64_t>(int64_t(offset)*65535/std::max<LONG>(1,length-1),0,65535));
}
POINT edgePoint(RECT r,int side,int position,int inset)
{
    position=std::clamp(position,0,65535);
    LONG x=r.left+LONG(int64_t(position)*std::max<LONG>(0,r.right-r.left-1)/65535);
    LONG y=r.top+LONG(int64_t(position)*std::max<LONG>(0,r.bottom-r.top-1)/65535);
    POINT p=side==0?POINT{r.left+inset,y}:side==1?POINT{r.right-1-inset,y}:side==2?POINT{x,r.top+inset}:POINT{x,r.bottom-1-inset};
    p.x=std::clamp(p.x,r.left,r.right-1);p.y=std::clamp(p.y,r.top,r.bottom-1);return p;
}
InputRouter::InputRouter(InputEnvironment e,std::function<bool(InputFrame)> s,int side,bool prefer,bool multiple)
    :env(std::move(e)),send(std::move(s)),edge(side),preferLocal(prefer),multi(multiple)
{if(side<0||side>3)throw std::runtime_error("Invalid screen position.");}
InputRouter::~InputRouter(){cancel(GetTickCount64());release();}
void InputRouter::mode(ControlMode value){if(mode_!=value){mode_=value;++revision_;if(env.modeChanged&&!(value==ControlMode::Local&&returning)&&!(value==ControlMode::Requesting&&transit))env.modeChanged(value);}}
bool InputRouter::emit(INPUT i)
{
    if(i.type==INPUT_KEYBOARD)i.ki.dwExtraInfo=remoteInputMarker();
    else i.mi.dwExtraInfo=remoteInputMarker();
    return env.desktop()&&env.inject(i);
}
bool InputRouter::move(POINT p,bool warp)
{
    auto r=env.bounds();if(r.right<=r.left||r.bottom<=r.top)return false;
    INPUT input{};input.type=INPUT_MOUSE;input.mi.dwFlags=MOUSEEVENTF_MOVE|MOUSEEVENTF_ABSOLUTE|MOUSEEVENTF_VIRTUALDESK;
    // Target pixel centres, avoiding a one-pixel bias on every relative mouse event.
    input.mi.dx=LONG((2*int64_t(std::clamp(p.x,r.left,r.right-1)-r.left)+1)*32768/(r.right-r.left));
    input.mi.dy=LONG((2*int64_t(std::clamp(p.y,r.top,r.bottom-1)-r.top)+1)*32768/(r.bottom-r.top));
    input.mi.dwExtraInfo=warp?cursorWarpMarker():remoteInputMarker();
    if(warp&&mode_==ControlMode::Sending)awaitingWarp=true;
    bool ok=env.desktop()&&env.inject(input);
    if(!ok)awaitingWarp=false;
    return ok;
}
bool InputRouter::post(InputFrame f){return connected&&send(f);}
bool InputRouter::held()const
{
    return std::any_of(physical.begin(),physical.end(),[](bool b){return b;})||env.held();
}
void InputRouter::release()
{
    for(size_t i=0;i<remoteKeys.size();++i)if(remoteKeys[i]&&emit(releases[i]))remoteKeys[i]=false;
    for(unsigned i=0;i<5;++i)if(remoteButtons[i]&&emit(button(i,false)))remoteButtons[i]=false;
}
void InputRouter::local(uint64_t now,int position)
{
    bool sending=mode_==ControlMode::Sending||transit;
    if(position<0)position=edgePosition(returnPoint,env.bounds(),edge);
    if(sending){returning=true;returnTarget=edgePoint(env.bounds(),edge,position);returnStarted=now;}
    release();mode(ControlMode::Local);cooldown=now+600;token=0;transit=false;crossEdge=-1;remainderX=remainderY=0;awaitingWarp=awaitingEntry=false;
    if(sending&&!move(returnTarget,true)){returning=false;if(env.modeChanged)env.modeChanged(ControlMode::Local);}
}
bool InputRouter::handoff(int side,int position,uint64_t now)
{
    if(!multi||mode_!=ControlMode::Sending||held()||side<0||side>3)return false;
    edge=side;transit=true;token=ownershipToken();
    requested=now;mode(ControlMode::Requesting);
    if(!post({InputType::Enter,token,edge,std::clamp(position,0,65535)})){cancel(now,"next computer unavailable");return false;}
    return true;
}
void InputRouter::returnAt(int side,int position,uint64_t now)
{
    if(side>=0&&side<4)edge=side;
    local(now,std::clamp(position,0,65535));
}
void InputRouter::cancel(uint64_t now,const char *reason)
{
    if(mode_!=ControlMode::Local)log(L"LAN input: "+wide(reason)+L"; mode="+std::to_wstring(int(mode_)));
    if(mode_!=ControlMode::Local)post({InputType::Return,token,edgePosition(env.cursor(),env.bounds(),edge)});
    local(now);
}
void InputRouter::tick(bool ready,uint64_t now)
{
    if(returning)
    {
        auto p=env.cursor();
        if((p.x==returnTarget.x&&p.y==returnTarget.y)||now-returnStarted>250)
        {returning=false;if(env.modeChanged)env.modeChanged(ControlMode::Local);}
    }
    connected=ready;
    if(!ready)cancel(now,"channel unavailable");
    else if(!env.desktop())cancel(now,"input desktop unavailable");
    else if(mode_==ControlMode::Requesting&&now-requested>1000)cancel(now,"handoff acknowledgement timed out");
    if(mode_==ControlMode::Sending&&awaitingWarp)
    {
        auto p=env.cursor();
        if(p.x==anchor.x&&p.y==anchor.y)awaitingWarp=false;
        else if(now-warpStarted>250)cancel(now,"source cursor parking did not complete");
    }
    if(mode_==ControlMode::Receiving&&awaitingEntry)
    {
        auto p=env.cursor();
        if(p.x==receivedPoint.x&&p.y==receivedPoint.y)
        {
            awaitingEntry=false;
            if(!post({InputType::Accept,token}))cancel(now,"cannot acknowledge cursor entry");
        }
        else if(now-entryStarted>250)cancel(now,"destination cursor entry did not complete");
    }
    if(mode_==ControlMode::Local)release(); // Retry releases blocked by a temporarily secure desktop.
}
void InputRouter::receive(const InputFrame &f,uint64_t now)
{
    if(!connected||!env.desktop()){cancel(now);return;}
    if(f.type==InputType::Enter)
    {
        if(!f.token||f.a<0||f.a>3||(!multi&&f.a!=(edge^1))||f.b<0||f.b>65535||held()||returning||
           mode_==ControlMode::Sending||mode_==ControlMode::Receiving||(mode_==ControlMode::Requesting&&preferLocal))
        {post({InputType::Return,f.token,32768});return;}
        if(mode_==ControlMode::Requesting)cancel(now);
        if(multi)edge=f.a^1;
        crossEdge=-1;token=f.token;remainderX=remainderY=0;mode(ControlMode::Receiving);
        receivedPoint=edgePoint(env.bounds(),edge,f.b);
        awaitingEntry=true;entryStarted=now;
        if(!move(receivedPoint))cancel(now,"destination cursor injection failed");
        return;
    }
    if(f.type==InputType::Accept)
    {
        if(mode_!=ControlMode::Requesting||f.token!=token||held()||(!transit&&!atEdge(env.cursor(),env.bounds(),edge))){post({InputType::Return,f.token,32768});if(f.token==token)cancel(now);return;}
        transit=false;auto r=env.bounds();anchor=env.anchor?env.anchor():POINT{(r.left+r.right)/2,(r.top+r.bottom)/2};mode(ControlMode::Sending);
        warpStarted=now;
        if(!move(anchor,true))cancel(now,"source cursor parking failed");
        return;
    }
    if(f.token!=token||!token)return;
    if(f.type==InputType::Return){log(L"LAN input: peer returned control");local(now,std::clamp(f.a,0,65535));return;}
    if(mode_!=ControlMode::Receiving)return;
    bool ok=true;
    if(f.type==InputType::Move)
    {
        if(std::abs(int64_t(f.a))>32768||std::abs(int64_t(f.b))>32768){cancel(now);return;}
        if(f.c!=0&&(f.c<48||f.c>768)){cancel(now);return;}
        // SendInput queues movement; GetCursorPos may still describe the old
        // centre while several network frames are drained in one timer tick.
        auto p=receivedPoint;auto r=env.bounds();
        double scale=f.c&&env.dpi?double(std::clamp(env.dpi(p),48,768))/f.c:1.0;
        double x=f.a*scale+remainderX,y=f.b*scale+remainderY;
        auto dx=int64_t(x),dy=int64_t(y);remainderX=x-dx;remainderY=y-dy;
        p.x=LONG(std::clamp<int64_t>(int64_t(p.x)+dx,r.left,r.right-1));p.y=LONG(std::clamp<int64_t>(int64_t(p.y)+dy,r.top,r.bottom-1));
        if(!multi&&atEdge(p,r,edge)&&!std::any_of(remoteKeys.begin(),remoteKeys.end(),[](bool b){return b;})&&
           !std::any_of(remoteButtons.begin(),remoteButtons.end(),[](bool b){return b;}))
        {log(L"LAN input: cursor crossed return edge");post({InputType::Return,token,edgePosition(p,r,edge)});local(now);return;}
        ok=move(p);if(ok)receivedPoint=p;
        if(multi&&ok&&!std::any_of(remoteKeys.begin(),remoteKeys.end(),[](bool b){return b;})&&
           !std::any_of(remoteButtons.begin(),remoteButtons.end(),[](bool b){return b;})){
            int horizontal=dx<0?0:1,vertical=dy<0?2:3,side=-1;
            if(std::abs(dy)>std::abs(dx)&&dy&&atEdge(p,r,vertical))side=vertical;
            else if(dx&&atEdge(p,r,horizontal))side=horizontal;
            else if(dy&&atEdge(p,r,vertical))side=vertical;
            if(side>=0&&(side!=crossEdge||now-crossAt>=150)){
                crossEdge=side;crossAt=now;ok=post({InputType::Cross,token,side,edgePosition(p,r,side)});
            }else if(side<0)crossEdge=-1;
        }
    }
    else if(f.type==InputType::Key)
    {
        if(f.a<1||f.a>255||f.b<0||f.b>255||f.c<0||f.c>3){cancel(now);return;}
        INPUT input{};input.type=INPUT_KEYBOARD;input.ki.wVk=WORD(f.a);input.ki.wScan=WORD(f.b);
        input.ki.dwFlags=(f.c&1?0:KEYEVENTF_KEYUP)|(f.c&2?KEYEVENTF_EXTENDEDKEY:0);
        bool modifier=f.a==VK_SHIFT||f.a==VK_CONTROL||f.a==VK_MENU||
            (f.a>=VK_LSHIFT&&f.a<=VK_RMENU)||f.a==VK_LWIN||f.a==VK_RWIN;
        // Explicit modifier VKs preserve their Windows state independently of
        // the sending keyboard's scan-code mapping. Text keys remain physical.
        if(f.b&&f.a!=VK_PAUSE&&!modifier)input.ki.dwFlags|=KEYEVENTF_SCANCODE;
        ok=emit(input);if(ok){remoteKeys[size_t(f.a)]=(f.c&1)!=0;input.ki.dwFlags|=KEYEVENTF_KEYUP;releases[size_t(f.a)]=input;}
    }
    else if(f.type==InputType::Button)
    {
        if(f.a<0||f.a>4||f.b<0||f.b>1){cancel(now);return;}
        ok=emit(button(unsigned(f.a),f.b!=0));if(ok)remoteButtons[size_t(f.a)]=f.b!=0;
    }
    else if(f.type==InputType::Wheel)
    {
        if(f.a<-32768||f.a>32767||f.b<0||f.b>1){cancel(now);return;}
        INPUT input{};input.type=INPUT_MOUSE;input.mi.dwFlags=f.b?MOUSEEVENTF_HWHEEL:MOUSEEVENTF_WHEEL;input.mi.mouseData=DWORD(f.a);ok=emit(input);
    }
    if(!ok)cancel(now,"Windows input injection failed");
}
bool InputRouter::keyboard(const KBDLLHOOKSTRUCT &k,bool down,uint64_t now)
{
    // Conversion/layout changes can produce synthetic modifier releases.
    // Never reinterpret our own tagged input as a physical takeover, even if
    // another hook has changed the injection flag.
    if(k.flags&LLKHF_INJECTED||k.dwExtraInfo==inputMarker()||
       k.dwExtraInfo==remoteInputMarker()||k.vkCode>255||!k.vkCode)return false;
    auto key=size_t(k.vkCode);physical[key]=down;
    if(mode_!=ControlMode::Local&&down&&key==VK_ESCAPE&&(physical[VK_LCONTROL]||physical[VK_RCONTROL]||physical[VK_CONTROL])&&
       (physical[VK_LMENU]||physical[VK_RMENU]||physical[VK_MENU]))
    {outKeys[key]=true;cancel(now);return true;}
    // A real local press takes over. A key-up alone is not evidence of intent:
    // Windows may generate it while changing the target keyboard layout.
    if(down&&(mode_==ControlMode::Receiving||mode_==ControlMode::Requesting))cancel(now,"local physical keyboard takeover");
    if(mode_==ControlMode::Sending)
    {
        outKeys[key]=down;
        if(!post({InputType::Key,token,int32_t(k.vkCode),int32_t(k.scanCode),int32_t((down?1:0)|((k.flags&LLKHF_EXTENDED)?2:0))}))cancel(now);
        return true;
    }
    if(outKeys[key]){if(!down)outKeys[key]=false;return true;}
    return false;
}
bool InputRouter::mouse(UINT message,const MSLLHOOKSTRUCT &m,uint64_t now,bool allowEnter)
{
    if((m.flags&LLMHF_INJECTED)&&m.dwExtraInfo==cursorWarpMarker())
    {awaitingWarp=false;return false;}
    if(m.flags&LLMHF_INJECTED)return false;
    if(returning&&message==WM_MOUSEMOVE)return true;
    int index=-1;bool down=false;
    switch(message){case WM_LBUTTONDOWN:down=true;[[fallthrough]];case WM_LBUTTONUP:index=0;break;
        case WM_RBUTTONDOWN:down=true;[[fallthrough]];case WM_RBUTTONUP:index=1;break;
        case WM_MBUTTONDOWN:down=true;[[fallthrough]];case WM_MBUTTONUP:index=2;break;
        case WM_XBUTTONDOWN:down=true;[[fallthrough]];case WM_XBUTTONUP:index=HIWORD(m.mouseData)==XBUTTON1?3:4;break;}
    if(mode_==ControlMode::Receiving&&message==WM_MOUSEMOVE)
    {
        if(m.pt.x==receivedPoint.x&&m.pt.y==receivedPoint.y)return false;
    }
    if(mode_==ControlMode::Receiving||(mode_==ControlMode::Requesting&&message!=WM_MOUSEMOVE))cancel(now,"local physical mouse takeover");
    if(mode_==ControlMode::Requesting&&!transit&&message==WM_MOUSEMOVE&&!atEdge(m.pt,env.bounds(),edge))cancel(now);
    if(mode_==ControlMode::Requesting&&transit&&message==WM_MOUSEMOVE)return true;
    if(mode_==ControlMode::Sending)
    {
        bool ok=true;
        if(message==WM_MOUSEMOVE)
        {
            // Do not interpret pre-warp edge coordinates as a huge delta from
            // the centre before Windows has actually parked the cursor.
            if(awaitingWarp)return true;
            auto dx=m.pt.x-anchor.x,dy=m.pt.y-anchor.y;
            if(dx||dy)ok=post({InputType::Move,token,dx,dy,env.dpi?std::clamp(env.dpi(anchor),48,768):0});
            // The hook consumes this physical motion, so Windows keeps the
            // cursor parked. Re-injecting an identical warp here can yield no
            // notification and must never become a gate for the next movement.
        }
        else if(index>=0){outButtons[size_t(index)]=down;ok=post({InputType::Button,token,index,down?1:0});}
        else if(message==WM_MOUSEWHEEL||message==WM_MOUSEHWHEEL)ok=post({InputType::Wheel,token,int16_t(HIWORD(m.mouseData)),message==WM_MOUSEHWHEEL?1:0});
        if(!ok)cancel(now);
        return true;
    }
    if(index>=0&&outButtons[size_t(index)]){if(!down)outButtons[size_t(index)]=false;return true;}
    if(mode_==ControlMode::Local&&connected&&allowEnter&&message==WM_MOUSEMOVE&&now>=cooldown&&
       env.desktop()&&!held()&&atEdge(m.pt,env.bounds(),edge))
    {
        token=ownershipToken();
        returnPoint=m.pt;requested=now;mode(ControlMode::Requesting);
        if(!post({InputType::Enter,token,edge,edgePosition(m.pt,env.bounds(),edge)}))cancel(now);
    }
    return false;
}
}