cli.py

import subprocess
from typing import List, Tuple
from kitty.boss import Boss
from kittens.tui.loop import debug
import sys

import json
import re

def parse_json_string(s):
    def try_parse_json(s):
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            return None
    
    # Try parsing the whole string
    parsed_json = try_parse_json(s)
    if parsed_json is not None:
        return parsed_json

    # Look for a code block
    code_block_pattern = r'```(?:json)?\s*([\s\S]*?)```'
    match = re.search(code_block_pattern, s, re.MULTILINE)

    if match:
        # Extract the code block and try parsing it
        code_block = match.group(1)
        parsed_json = try_parse_json(code_block)
        if parsed_json is not None:
            return parsed_json

    return None 

def main(args: List[str]) -> str:
    return sys.stdin.read()

def split_buffer(s: str) -> Tuple[str, str]:
    buffer = ""
    request = None

    if s.strip():
        lines = s.split('\n')
        last_line = lines[-1]

        if '#' in last_line:
            buffer = '\n'.join(lines[:-1])
            request = last_line.split('#', 1)[1].strip()
        else:
            buffer = s
            request = None
        
        return buffer, request
    else:
        return s, None

def run_pinocchio(buffer, request):
    command = ["pinocchio", "mine", "zsh"]
    if request is not None:
        command += ["--request", request]
    command += ["-"]

    process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    stdout, stderr = process.communicate(buffer)

    if process.returncode != 0:
        return None

    return stdout

def gum_choose(suggestions) -> str:
    commands = [s["command"] for s in suggestions]
    command = ["gum", "choose"] + commands

    process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    stdout, stderr = process.communicate()

    if process.returncode != 0:
        return None
    else:
        return stdout

from kittens.tui.handler import result_handler
@result_handler(type_of_input='screen')
def handle_result(args: List[str], answer: str, target_window_id: int, boss: Boss) -> None:
    debug("answer: " + answer + "\n")
    buffer, request = split_buffer(answer)
    debug("buffer: " + buffer + "\n")
    if request is None:
        debug("no request\n")
    else:
        debug("request: " + request + "\n")
    suggestions = run_pinocchio(buffer, request)
    debug("suggestions: " + suggestions + "\n")
    json_string = parse_json_string(suggestions)
    commands = []
    selection = "-- Could not get suggestions\n"
    if json_string is None:
        debug("no json\n")
    else:  
        debug("json: " + json.dumps(json_string) + "\n")
        selection = "\n$(gum choose " + " ".join(["\"" + s["command"] + "\"" for s in json_string]) + ")"
        # choice = gum_choose(json_string)
        # if choice is not None:
        #     selection = '\n' + choice

    # get the kitty window into which to paste answer
    w = boss.window_id_map.get(target_window_id)
    if w is not None:
        w.paste_text(selection)