import times, tables, strformat, strutils, algorithm
import imguin/[cimgui, glfw_opengl, simple]

import ./console
import ../core/[task, websocket]
import ../utils/[appImGui, colors]
import ../../modules/manager
import ../../common/types

type 
    SessionsTableComponent* = ref object of RootObj
        title: string 
        agents*: seq[UIAgent]
        agentActivity*: Table[string, int64]                # Direct O(1) access to latest checkin
        agentImpersonation*: Table[string, string]
        selection: ptr ImGuiSelectionBasicStorage
        consoles: ptr Table[string, ConsoleComponent]
        focusedConsole*: string

proc SessionsTable*(title: string, consoles: ptr Table[string, ConsoleComponent]): SessionsTableComponent = 
    result = new SessionsTableComponent
    result.title = title
    result.agents = @[]
    result.agentActivity = initTable[string, int64]()
    result.selection = ImGuiSelectionBasicStorage_ImGuiSelectionBasicStorage()
    result.consoles = consoles
    result.focusedConsole = ""

proc cmp(x, y: UIAgent): int =
    return cmp(x.firstCheckin, y.firstCheckin)

proc interact(component: SessionsTableComponent) = 
    # Open a new console for each selected agent session
    var it: pointer = nil
    var row: ImGuiID

    while ImGuiSelectionBasicStorage_GetNextSelectedItem(component.selection, addr it, addr row):
        let agent = component.agents[cast[int](row)]

        # Create a new console window
        if not component.consoles[].hasKey(agent.agentId):
            component.consoles[][agent.agentId] = Console(agent)

        component.focusedConsole = fmt"[{agent.agentId}] {agent.username}@{agent.hostname}"
    
    component.selection.ImGuiSelectionBasicStorage_Clear()

proc draw*(component: SessionsTableComponent, showComponent: ptr bool, connection: WsConnection) = 
    igBegin(component.title.cstring, showComponent, 0)

    let textSpacing = igGetStyle().ItemSpacing.x

    let tableFlags = (
        ImGuiTableFlags_Resizable.int32 or 
        ImGuiTableFlags_Reorderable.int32 or 
        ImGuiTableFlags_Hideable.int32 or 
        ImGuiTableFlags_HighlightHoveredColumn.int32 or 
        ImGuiTableFlags_RowBg.int32 or 
        ImGuiTableFlags_BordersV.int32 or 
        ImGuiTableFlags_BordersH.int32 or 
        ImGuiTableFlags_ScrollY.int32 or
        ImGuiTableFlags_ScrollX.int32 or 
        ImGuiTableFlags_NoBordersInBodyUntilResize.int32 or
        ImGui_TableFlags_SizingStretchSame.int32
    )

    let cols: int32 = 12
    if igBeginTable("Sessions", cols, tableFlags, vec2(0.0f, 0.0f), 0.0f):

        igTableSetupColumn("AgentID", ImGuiTableColumnFlags_NoReorder.int32 or ImGuiTableColumnFlags_NoHide.int32, 0.0f, 0)
        igTableSetupColumn("ListenerID", ImGuiTableColumnFlags_DefaultHide.int32, 0.0f, 0)
        igTableSetupColumn("IP (Internal)", ImGuiTableColumnFlags_None.int32, 0.0f, 0)
        igTableSetupColumn("IP (External)", ImGuiTableColumnFlags_DefaultHide.int32, 0.0f, 0)
        igTableSetupColumn("Username", ImGuiTableColumnFlags_None.int32, 0.0f, 0)
        igTableSetupColumn("Hostname", ImGuiTableColumnFlags_None.int32, 0.0f, 0)
        igTableSetupColumn("Domain", ImGuiTableColumnFlags_None.int32, 0.0f, 0)
        igTableSetupColumn("OS", ImGuiTableColumnFlags_None.int32, 0.0f, 0)
        igTableSetupColumn("Process", ImGuiTableColumnFlags_None.int32, 0.0f, 0)
        igTableSetupColumn("PID", ImGuiTableColumnFlags_None.int32, 0.0f, 0)
        igTableSetupColumn("First seen", ImGuiTableColumnFlags_None.int32, 0.0f, 0)
        igTableSetupColumn("Last seen", ImGuiTableColumnFlags_None.int32, 0.0f, 0)

        igTableSetupScrollFreeze(0, 1)
        igTableHeadersRow()
        
        var multiSelectIO = igBeginMultiSelect(ImGuiMultiSelectFlags_ClearOnEscape.int32 or ImGuiMultiSelectFlags_BoxSelect1d.int32, component.selection[].Size, int32(component.agents.len())) 
        ImGuiSelectionBasicStorage_ApplyRequests(component.selection, multiSelectIO)

        # Sort sessions table based on first checkin
        component.agents.sort(cmp)
        for row, agent in component.agents: 
            igTableNextRow(ImGuiTableRowFlags_None.int32, 0.0f)

            if igTableSetColumnIndex(0):          
                # Enable multi-select functionality       
                igSetNextItemSelectionUserData(row)
                var isSelected = ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](row))
                
                # Highlight high integrity sessions in red
                if agent.elevated: 
                    igPushStyleColor_Vec4(ImGui_Col_Text.cint, CONSOLE_ERROR)

                discard igSelectable_Bool(agent.agentId.cstring, isSelected, ImGuiSelectableFlags_SpanAllColumns.int32, vec2(0.0f, 0.0f))
                
                if agent.elevated:
                    igPopStyleColor(1)

                # Interact with session on double-click
                if igIsMouseDoubleClicked_Nil(ImGui_MouseButton_Left.int32):
                    component.interact()

            if igTableSetColumnIndex(1): 
                igText(agent.listenerId.cstring)
            if igTableSetColumnIndex(2): 
                igText(agent.ipInternal.cstring)
            if igTableSetColumnIndex(3): 
                igText(agent.ipExternal.cstring)
            if igTableSetColumnIndex(4): 

                igText(agent.username.cstring)
                if component.agentImpersonation.hasKey(agent.agentId):
                    igSameLine(0.0f, textSpacing)
                    igText(fmt"[{component.agentImpersonation[agent.agentId]}]".cstring)

            if igTableSetColumnIndex(5): 
                igText(agent.hostname.cstring)
            if igTableSetColumnIndex(6): 
                igText(agent.domain.cstring)
            if igTableSetColumnIndex(7): 
                igText(agent.os.cstring)
            if igTableSetColumnIndex(8): 
                igText(agent.process.cstring)
            if igTableSetColumnIndex(9): 
                igText(($agent.pid).cstring)
            if igTableSetColumnIndex(10): 
                let duration = now() - agent.firstCheckin.fromUnix().local()
                let totalSeconds = duration.inSeconds
                    
                let hours = totalSeconds div 3600
                let minutes = (totalSeconds mod 3600) div 60
                let seconds = totalSeconds mod 60
                
                igText(fmt"{hours:02d}:{minutes:02d}:{seconds:02d} ago".cstring)

            if igTableSetColumnIndex(11): 
                let duration = now() - component.agentActivity[agent.agentId].fromUnix().local()
                let totalSeconds = duration.inSeconds
                
                let hours = totalSeconds div 3600
                let minutes = (totalSeconds mod 3600) div 60
                let seconds = totalSeconds mod 60
                
                let timeText = fmt"{hours:02d}:{minutes:02d}:{seconds:02d} ago"
                if totalSeconds > agent.sleep: 
                    igTextColored(GRAY, timeText.cstring)
                else: 
                    igText(timeText.cstring)

        # Handle right-click context menu
        # Right-clicking the table header to hide/show columns or reset the layout is only possible when no sessions are selected
        if component.selection[].Size > 0 and igBeginPopupContextWindow("TableContextMenu", ImGui_PopupFlags_MouseButtonRight.int32): 

            if igMenuItem("Interact", nil, false, true): 
                component.interact()
                igCloseCurrentPopup()
        
            if igBeginMenu("Exit", true):
                if igMenuItem("Process", nil, false, true): 
                    for i, agent in component.agents:
                        if ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
                            if component.consoles[].hasKey(agent.agentId):
                                component.consoles[][agent.agentId].handleAgentCommand(connection, "exit process")
                            else: 
                                let task = createTask(agent.agentId, agent.listenerId, getCommandByType(CMD_EXIT), @["process"])
                                connection.sendAgentTask(agent.agentId, "exit process", task)

                    ImGuiSelectionBasicStorage_Clear(component.selection)
                    igCloseCurrentPopup()

                if igMenuItem("Thread", nil, false, true):
                    for i, agent in component.agents:
                        if ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
                            if component.consoles[].hasKey(agent.agentId):
                                component.consoles[][agent.agentId].handleAgentCommand(connection, "exit thread") 
                            else: 
                                let task = createTask(agent.agentId, agent.listenerId, getCommandByType(CMD_EXIT), @["thread"])
                                connection.sendAgentTask(agent.agentId, "exit thread", task)

                    ImGuiSelectionBasicStorage_Clear(component.selection)
                    igCloseCurrentPopup()

                igEndMenu()

            igSeparator()

            if igMenuItem("Remove", nil, false, true): 
                # Update agents table with only non-selected ones
                var newAgents: seq[UIAgent] = @[]
                for i, agent in component.agents:
                    if not ImGuiSelectionBasicStorage_Contains(component.selection, cast[ImGuiID](i)):
                        newAgents.add(agent)
                    else: 
                        # Send message to team server to remove delete the agent from the database and stop it from re-appearing when the client is restarted
                        connection.sendAgentRemove(agent.agentId)

                component.agents = newAgents
                ImGuiSelectionBasicStorage_Clear(component.selection)
                igCloseCurrentPopup()

            igEndPopup()

        multiSelectIO = igEndMultiSelect()
        ImGuiSelectionBasicStorage_ApplyRequests(component.selection, multiSelectIO)
        
        # Auto-scroll to bottom
        if igGetScrollY() >= igGetScrollMaxY():
            igSetScrollHereY(1.0f)

        igEndTable()

    igEnd()