PoshCodex.psm1

### --- PUBLIC FUNCTIONS --- ###
#Region - Enter-CompletionKeybind.ps1
#Requires -Modules PSReadLine

# Asks for user input (reads key combinations), calls internal functions to convert to string and set it as the new keybind
function Enter-CompletionKeybind {
    # Add cmdletBinding to the parameter list
    [CmdletBinding()]
    param()

    Write-Host "Press any key or combination (Ctrl, Alt, Shift + other keys) to set the new keybind for PoshCodex. Press 'Escape' to exit."
    # Read the key press
    $key = [System.Console]::ReadKey($true)
    $new_autocomplete_keybind = Convert-KeyPressToString $key
    if ($null -eq $new_autocomplete_keybind) {
        # exit, no input given (user pressed Escape).
        return
    }

    # get old keybind, call set keybind for new
    $old_autocomplete_keybind = [Environment]::GetEnvironmentVariable('AUTOCOMPLETE_KEYBIND', 'User')
    Set-CompletionKeybind $old_autocomplete_keybind $new_autocomplete_keybind
    Write-Host "New keybind set: $new_autocomplete_keybind"

    # Update env var with new keybind
    [Environment]::SetEnvironmentVariable('AUTOCOMPLETE_KEYBIND', $new_autocomplete_keybind, [System.EnvironmentVariableTarget]::User)
    # Write-Host "Value of newly set env var: $([Environment]::GetEnvironmentVariable('AUTOCOMPLETE_KEYBIND', 'User'))"
}
Export-ModuleMember -Function Enter-CompletionKeybind
#EndRegion - Enter-CompletionKeybind.ps1
### --- PRIVATE FUNCTIONS --- ###
#Region - Convert-KeyPressToString.ps1
#Requires -Modules PSReadLine

# Function to convert a user input key press into a keybind string
function Convert-KeyPressToString {
    # Add cmdletBinding to the parameter list
    [CmdletBinding()]
    param($key)

    # First, check and append Ctrl if it's pressed
    if ($key.Modifiers -band [ConsoleModifiers]::Control) {
        $keybind += 'Ctrl+'
    }

    # Next, check and append Alt if it's pressed
    if ($key.Modifiers -band [ConsoleModifiers]::Alt) {
        $keybind += 'Alt+'
    }

    # Lastly, check and append Shift if it's pressed
    if ($key.Modifiers -band [ConsoleModifiers]::Shift) {
        $keybind += 'Shift+'
    }

    # Append the actual key that was pressed
    $keybind += $key.Key.ToString()

    # Display the keybind string
    # Write-Host "Keybind entered: $keybind"

    # If the Escape key is pressed, exit the loop
    if ($key.Key -eq 'Escape') {
        Write-Host 'Aborted by user, exiting...'
        return $null
    }

    return $keybind
}
#EndRegion - Convert-KeyPressToString.ps1
#Region - Invoke-Ollama-Api.ps1
function Invoke-Ollama-Api {
    [CmdletBinding()]
    param (
        $BUFFER
    )

    # [Microsoft.PowerShell.PSConsoleReadLine]::AddToHistory($line)
    # [Microsoft.PowerShell.PSConsoleReadLine]::RevertLine()

    $data = @{
        model  = "$env:OLLAMA_MODEL"
        prompt = $BUFFER
        stream = $false
    }

    $json_output = Invoke-RestMethod -Method POST `
        -Uri "$env:OLLAMA_HOST/api/generate" `
        -Body ($data | ConvertTo-Json) `
        -ContentType 'application/json; charset=utf-8';

    return $json_output
}
#EndRegion - Invoke-Ollama-Api.ps1
#Region - Set-CompletionKeybind.ps1
#Requires -Modules PSReadLine

# Function to simply handle the setting and removal of keybinds, without worrying about the user input.
function Set-CompletionKeybind {
    # Add cmdletBinding to the parameter list
    [CmdletBinding()]
    param(
        $old_keybind, $new_keybind
    )

    if ($null -ne $old_keybind) {
        # unset current handler for Write-Completion if it exists
        Remove-PSReadLineKeyHandler -Chord $old_keybind
        Write-Host "Previous keybind removed: $old_keybind"
    }
    
    Set-PSReadLineKeyHandler -Chord $new_keybind `
        -BriefDescription Write-Completion `
        -LongDescription 'Autocomplete the command' `
        -ScriptBlock { Write-Completion }
}
#EndRegion - Set-CompletionKeybind.ps1
#Region - Write-Completion.ps1
#Requires -Modules PSReadLine

function Write-Completion {
    # Add cmdletBinding to the parameter list
    [CmdletBinding()]
    param()

    $BUFFER = $null
    $cursor = $null

    # read text from current buffer
    [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$BUFFER, [ref]$cursor)
    
    # If the buffer text itself contains double quotes, then we need to escape them.
    $BUFFER = $BUFFER.Replace('"', '""')

    $json_output = Invoke-Ollama-Api $BUFFER

    # check if json_output is not equal to null
    if ($null -ne $json_output) {
        $completion = $json_output.response

        # Insert the completion on the next line. This will NOT cause the command to be executed.
        [Microsoft.PowerShell.PSConsoleReadLine]::InsertLineBelow();
        [Microsoft.PowerShell.PSConsoleReadLine]::Insert($completion)
    }
    else {
        Write-Output 'Response returned by API is null! It could be an internal error or the model is not installed properly hrough Ollama. Please fix and try again.' -ForegroundColor Red
    }
}
#EndRegion - Write-Completion.ps1
## Set necessary environment variables:

[System.Environment]::SetEnvironmentVariable('OLLAMA_HOST', 'http://localhost:11434', [System.EnvironmentVariableTarget]::User)
[System.Environment]::SetEnvironmentVariable('OLLAMA_MODEL', 'rishi255/posh_codex_model', [System.EnvironmentVariableTarget]::User)


## Set default keybind:

$default_keybind = 'Ctrl+Shift+O'
Set-CompletionKeybind $null $default_keybind;
[Environment]::SetEnvironmentVariable('AUTOCOMPLETE_KEYBIND', $default_keybind, [System.EnvironmentVariableTarget]::User)