#!/usr/bin/env python3

import re
import sys
import argparse
import os
import stat
from dataclasses import dataclass
from typing import List, Union

TOKEN_SPECIFICATION = [
    ('IMPORT', r'import'),
    ('EVAL', r'eval'),
    ('OPTION', r'option'),
    ('DMENU', r'dmenu'),
    ('LPAREN', r'\('),
    ('RPAREN', r'\)'),
    ('STRING', r'"(?:\\.|[^"\\])*"|\'[^\']*\'|`[^`]*`'),
    ('SKIP', r'[\n \t]+'),
    ('MISMATCH', r'.'),
]

TOKEN_REGEX = '|'.join(f'(?P<{name}>{pattern})' for name, pattern in TOKEN_SPECIFICATION)


@dataclass
class Token:
    type: str
    value: str
    pos: int

    def __repr__(self):
        return f'Token({self.type}, {self.value}, pos={self.pos})'


def tokenize(code):
    tokens = []
    for mo in re.finditer(TOKEN_REGEX, code):
        kind = mo.lastgroup
        value = mo.group()
        pos = mo.start()
        if kind == 'SKIP':
            continue
        elif kind == 'MISMATCH':
            raise RuntimeError(f'Unexpected token {value!r} at position {pos}')
        else:
            tokens.append(Token(kind, value, pos))
    return tokens


def restore_string(string_tkn_val: str):
    if string_tkn_val[0] == '"':
        return string_tkn_val[1:-1].replace("\\\"", "\"").replace("\\\\", "\\")
    else:
        return string_tkn_val[1:-1]


class ActionDmenu:
    ...


@dataclass
class OptionInDmenu:
    label: str
    action: Union[ActionDmenu, str]

    def __repr__(self):
        if isinstance(self.action, ActionDmenu):
            return f"option {self.label} {self.action}"
        else:
            return f"option {self.label} eval {self.action}"


@dataclass
class ActionDmenu:
    prompt: str
    options: List[OptionInDmenu]

    def __repr__(self):
        return f"dmenu {self.prompt} ( {' ; '.join(map(str, self.options))} )"


class Parser:
    def __init__(self, tokens):
        self.tokens = tokens
        self.pos = 0

    def current(self):
        if self.pos < len(self.tokens):
            return self.tokens[self.pos]
        return None

    def eat(self, token_type) -> Token:
        token = self.current()
        if token is None:
            raise RuntimeError("Unexpected end of input")
        if token.type == token_type:
            self.pos += 1
            return token
        raise RuntimeError(f"Expected token {token_type} at position {token.pos} but got {token.type}")

    # Grammar production: <option> ::= option <string> ( eval <string> | <dmenu> )
    def parse_option(self):
        self.eat('OPTION')
        label_tkn = self.eat('STRING')
        label = restore_string(label_tkn.value)
        if '\t' in label or '\n' in label or '\000' in label:
            raise RuntimeError(f"Label of option at {label_tkn.pos} contains illegal character")
        current = self.current()
        if current is None:
            raise RuntimeError("Expected end-of-file after option")
        if current.type == 'EVAL':
            self.eat('EVAL')
            action_command_tkn = self.eat('STRING').value
            action = restore_string(action_command_tkn)
            if '\n' in action or '\000' in action:
                raise RuntimeError(f"Command at {action_command_tkn.pos} contains illegal character")
        elif current.type == 'DMENU':
            action = self.parse_dmenu()
        else:
            raise RuntimeError(f"Unexpected token {current.type} in option at position {current.pos}")
        return OptionInDmenu(label, action)

    # Grammar production: <dmenu> ::= dmenu <string> \( <option> * \)
    def parse_dmenu(self):
        self.eat('DMENU')
        prompt = restore_string(self.eat('STRING').value)
        self.eat('LPAREN')
        options = []
        while self.current() is not None and self.current().type != 'RPAREN':
            options.append(self.parse_option())
        self.eat('RPAREN')
        return ActionDmenu(prompt, options)

    def parse_file_as_dmenu(self):
        node = self.parse_dmenu()
        if self.current() is not None:
            raise RuntimeError("Extra tokens after valid dmenu")
        return node


def shell_escape(s: str) -> str:
    return "'" + s.replace("'", "'\\''") + "'"


def write_to_file(path: str, text: str) -> None:
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)

def make_file_executable(path: str):
    old = os.stat(path).st_mode
    os.chmod(path, old | stat.S_IXUSR | stat.S_IXGRP)


def generate_tsv_file(dpath: str, cur_tsv_name, options: List[OptionInDmenu]):
    def option_to_line(i: int) -> str:
        opt = options[i]
        if isinstance(opt.action, ActionDmenu):
            child_tsv_name = f"{cur_tsv_name}-{i}"
            generate_tsv_file(dpath, child_tsv_name, opt.action.options)
            return f"{opt.label}\t<{shell_escape(os.path.join(dpath, child_tsv_name))} dmenu " + \
                f"-p {shell_escape(opt.action.prompt)} | tail -n 1 | sh"
        else:
            return f"{opt.label}\t{opt.action}"

    write_to_file(
        os.path.join(dpath, cur_tsv_name),
        "\n".join(map(option_to_line, range(len(options))))
    )


def compile_dmenuctree(dpath: str, parsed_tree: ActionDmenu):
    print(f"dmenu tabfiles root is {dpath}")
    main_script = os.path.join(dpath, "invoke.sh")
    cur_tsv_name = "tsv"
    invoke_snippet = f"#!/bin/sh\n" + \
                     f"<{shell_escape(os.path.join(dpath, cur_tsv_name))} dmenu -p {shell_escape(parsed_tree.prompt)} " + \
                     f"| tail -n 1 | sh"
    write_to_file(main_script, invoke_snippet)
    make_file_executable(main_script)
    generate_tsv_file(dpath, cur_tsv_name, parsed_tree.options)


def main():
    parser_arg = argparse.ArgumentParser(description="Parse a file using custom dmenu syntax")
    parser_arg.add_argument("input_file", help="Path to the input file")
    parser_arg.add_argument("dest_dir", help="Destination directory path")
    args = parser_arg.parse_args()

    if not os.path.exists(args.input_file):
        print("Input file does not exist.", file=sys.stderr)
        sys.exit(1)
    if not os.path.isdir(args.dest_dir):
        os.mkdir(args.dest_dir)

    with open(args.input_file, 'r') as f:
        code = f.read()

    try:
        tokens = tokenize(code)
        parsed_tree = Parser(tokens).parse_file_as_dmenu()
    except Exception as e:
        print(f"Error during parsing: {e}", file=sys.stderr)
        sys.exit(1)

    compile_dmenuctree(os.path.realpath(args.dest_dir), parsed_tree)


if __name__ == '__main__':
    main()