#!/usr/bin/env python3

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

def find_line_column(code, pos):
    last_newline = code.rfind('\n', 0, pos)
    line = code.count('\n', 0, pos)
    if last_newline == -1:
        column = pos
    else:
        column = pos - last_newline - 1
    return (line, column)

def stringify_code_pos(code, pos):
    line, column = find_line_column(code, pos)
    return f"{line+1}:{column+1}"

TOKEN_SPECIFICATION = [
    ('IMPORT', r'import'),
    ('OPTIONS', r'options'),
    ('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, filename):
    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 {stringify_code_pos(code, pos)}. File {filename}')
        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))} )"


def resolve_path_relative_to_file(path_a, path_b):
    if os.path.isabs(path_b):
        return path_b
    a_dir = os.path.dirname(os.path.abspath(path_a))
    path_c = os.path.normpath(os.path.join(a_dir, path_b))
    return path_c


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

    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 {stringify_code_pos(self.code, self.pos)} but got {token.type}. File {self.filename}")

    # Grammar production: <option> ::= option <string> ( eval <string> | <dmenu> ) | import options <string (filepath)>
    def parse_option_and_append(self, targ_arr):
        if self.current().type != None and self.current().type == 'OPTION':
            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 {stringify_code_pos(self.code, label_tkn.pos)} contains illegal character. File {self.filename}")
            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')
                action = restore_string(action_command_tkn.value)
                if '\n' in action or '\000' in action:
                    raise RuntimeError(f"Command at {stringify_code_pos(self.code, action_command_tkn.pos)} contains illegal character. File {self.filename}")
            elif current.type == 'DMENU':
                action = self.parse_dmenu()
            else:
                raise RuntimeError(f"Unexpected token {current.type} in option at position {stringify_code_pos(self.code, current.pos)}. File {self.filename}")
            targ_arr.append(OptionInDmenu(label, action))
        else:
            self.eat('IMPORT')
            self.eat('OPTIONS')
            filepath_tkn = self.eat('STRING')
            imported_file_path = restore_string(filepath_tkn.value)
            child_file = resolve_path_relative_to_file(self.filename, imported_file_path)
                
            print(f"Importing dmenul file: {child_file}")
            with open(child_file, 'r') as f:
                child_file_code = f.read()

            child_file_tokens = tokenize(child_file_code, self.filename)
            parsed_child = Parser(child_file_tokens, child_file, child_file_code).parse_file_as_dmenul()

            targ_arr += parsed_child


    # 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':
            self.parse_option_and_append(options)
        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 root. File {self.filename}")
        return node

    def parse_file_as_dmenul(self):
        options = []
        while self.current() is not None:
            self.parse_option_and_append(options)
        return options


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 -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, init_shell_file_path: str):
    print(f"dmenu tabfiles root is {dpath}")
    main_script = os.path.join(dpath, "invoke.sh")
    cur_tsv_name = "tsv"
    invoke_snippet = "#!/bin/sh\n"
    if len(init_shell_file_path) > 0:
        invoke_snippet += f"[ -f {shell_escape(init_shell_file_path)} ] && source {shell_escape(init_shell_file_path)} \n"
    invoke_snippet += f"<{shell_escape(os.path.join(dpath, cur_tsv_name))} dmenu -p {shell_escape(parsed_tree.prompt)} | tail -n 1 | sh\n"
    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="Compile .dmenu file into directory of executable scripts")
    parser_arg.add_argument("input_file", help="Path to the input file")
    parser_arg.add_argument("dest_dir", help="Destination directory path")
    parser_arg.add_argument("init_shell_file_path", help="Invoke.sh script will source this shell script before running this dmenu conglomeration")
    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)
    init_shell_file_path = args.init_shell_file_path if args.init_shell_file_path else ""

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

    try:
        tokens = tokenize(code, args.input_file)
        parsed_tree = Parser(tokens, args.input_file, code).parse_file_as_dmenu()
    except Exception as e:
        print(f"Parsing error: {e}", file=sys.stderr)
        sys.exit(1)

    compile_dmenuctree(os.path.realpath(args.dest_dir), parsed_tree, resolve_path_relative_to_file(args.input_file, init_shell_file_path))


if __name__ == '__main__':
    main()