import re
import yaml

from abc import abstractmethod
from dataclasses import dataclass
from typing import Dict, List, Optional
from simple_parsing.helpers import FrozenSerializable


@dataclass(frozen=True)
class AssistantMetadata(FrozenSerializable):
    """Pass observations to the assistant, and get back a response."""
    system_template: Optional[str] = None
    instance_template: Optional[str] = None


# TODO: first can be used for two-stage actions
# TODO: eventually might control high-level control flow
@dataclass(frozen=True)
class ControlMetadata(FrozenSerializable):
    """TODO: should be able to control high-level control flow after calling this command"""
    next_step_template: Optional[str] = None
    next_step_action_template: Optional[str] = None
    

@dataclass(frozen=True)
class Command(FrozenSerializable):
    code: str
    name: str
    docstring: Optional[str] = None
    end_name: Optional[str] = None  # if there is an end_name, then it is a multi-line command
    arguments: Optional[Dict] = None
    signature: Optional[str] = None


class ParseCommandMeta(type):
    _registry = {}

    def __new__(cls, name, bases, attrs):
        new_cls = super().__new__(cls, name, bases, attrs)
        if name != "ParseCommand":
            cls._registry[name] = new_cls
        return new_cls


@dataclass
class ParseCommand(metaclass=ParseCommandMeta):
    @classmethod
    def get(cls, name):
        try:
            return cls._registry[name]()
        except KeyError:
            raise ValueError(f"Command parser ({name}) not found.")

    @abstractmethod
    def parse_command_file(self, path: str) -> List[Command]:
        """
        Define how to parse a file into a list of commands.
        """
        raise NotImplementedError
    
    @abstractmethod
    def generate_command_docs(self, commands: List[Command], subroutine_types, **kwargs) -> str:
        """
        Generate a string of documentation for the given commands and subroutine types.
        """
        raise NotImplementedError


# DEFINE NEW COMMAND PARSER FUNCTIONS BELOW THIS LINE

class ParseCommandBash(ParseCommand):
    def parse_command_file(self, path: str) -> List[Command]:
        """
        Simple logic for parsing a bash file and segmenting it into functions.

        Assumes that all functions have their name and opening curly bracket in one line,
        and closing curly bracket in a line by itself.
        """
        with open(path, "r") as file:
            lines = file.readlines()

        commands = []

        comment = ""
        i = 0
        while i < len(lines):
            line = lines[i]
            i += 1
            if line.startswith("# "):
                comment += line[2:]
            elif line.strip().endswith("() {"):
                # Start of a Bash function
                name = line.split()[0][:-2]

                code = line
                while lines[i].strip() != "}":
                    code += lines[i]
                    i += 1
                code += lines[i]

                if comment.startswith("@yaml"):
                    # Parse the yaml docstring
                    kwargs = yaml.safe_load(comment.split("\n", 1)[1])
                else:
                    kwargs = dict(signature=name, docstring=comment.strip())

                kwargs["name"] = name
                kwargs["code"] = code

                commands.append(Command.from_dict(kwargs))
                comment = ""
            else:
                comment = ""

        return commands


    def generate_command_docs(self, commands: List[Command], subroutine_types, **kwargs) -> str:
        docs = ""
        for cmd in commands:
            if cmd.docstring is not None:
                docs += f"{cmd.signature or cmd.name} - {cmd.docstring.format(**kwargs)}\n"
        for subroutine in subroutine_types:
            if subroutine.docstring is not None:
                docs += f"{subroutine.signature or subroutine.name} - {subroutine.docstring.format(**kwargs)}\n"
        return docs


class ParseCommandPython(ParseCommand):
    def parse_command_file(self, path: str) -> List[Command]:
        with open(path, "r") as file:
            lines = file.readlines()
        commands = []
        idx = 0
        docs = []
        while idx < len(lines):
            line = lines[idx]
            idx += 1
            if line.startswith("# "):
                docs.append(line[2:])
            elif line.strip().endswith("() {"):
                name = line.split()[0][:-2]
                code = line
                while lines[idx].strip() != "}":
                    code += lines[idx]
                    idx += 1
                code += lines[idx]

                docstring, end_name, arguments, signature = None, None, None, name
                docs_dict = yaml.safe_load("".join(docs).replace('@yaml', ''))
                if docs_dict is not None:
                    docstring = docs_dict["docstring"]
                    end_name = docs_dict.get("end_name", None)
                    arguments = docs_dict["arguments"] if "arguments" in docs_dict else None

                    if "signature" in docs_dict:
                        signature = docs_dict["signature"]
                    else:
                        if arguments is not None:
                            for param, settings in arguments.items():
                                if settings["required"]:
                                    signature += f" <{param}>"
                                else:
                                    signature += f" [<{param}>]"

                command = Command.from_dict({
                    "code": code,
                    "docstring": docstring,
                    "end_name": end_name,
                    "name": name,
                    "arguments": arguments,
                    "signature": signature
                })
                commands.append(command)
                docs = []
        return commands


    def generate_command_docs(self, commands: List[Command], subroutine_types, **kwargs) -> str:
        def helper_old(cmd):
            indent = "  "
            docs = f"{cmd.signature}\n"
            if "arguments" in cmd.__dict__ and cmd.arguments is not None:
                docs += f"{indent}- Arguments:\n"
                for param, settings in cmd.arguments.items():
                    docs += f"{indent*2}- {param}"
                    if not settings["required"]:
                        docs += " (optional)"
                    docs += f" [{settings['type']}]: {settings['description']}\n"
            return docs
        
        def helper(cmd):
            docs = f"// {cmd.docstring}\n"
            docs += f"type {cmd.name} = ("
            if "arguments" in cmd.__dict__ and cmd.arguments is not None:
                docs += "_: {\n"
                for param, settings in cmd.arguments.items():
                    docs += f"// {settings['description']}\n"
                    if settings["required"]:
                        docs += f"{param}: {settings['type']},\n"
                    else:
                        docs += f"{param}?: {settings['type']},\n"
                docs += "}"
            docs += ") => any;\n"
            return docs

        docs = []
        for cmd in commands:
            if cmd.docstring is not None:
                docs.append(helper(cmd).strip())
                    
        for subroutine in subroutine_types:
            if subroutine.docstring is not None:
                docs.append(helper(subroutine).strip())

        docs = sorted(docs)
        docs = "\n\n".join(docs)
        return docs.strip()


class ParseCommandDetailed(ParseCommandPython):
    """
    # command_name:
    #   "docstring"
    #   signature: "signature"
    #   arguments:
    #     arg1 (type) [required]: "description"
    #     arg2 (type) [optional]: "description"
    """
    def get_signature(cmd):
        signature = cmd.name
        if "arguments" in cmd.__dict__ and cmd.arguments is not None:
                for param, settings in cmd.arguments.items():
                    if settings["required"]:
                        signature += f" <{param}>"
                    else:
                        signature += f" [<{param}>]"
        return signature

    def generate_command_docs(
            self,
            commands: List[Command],
            subroutine_types,
            **kwargs,
            ) -> str:
        docs = ""
        for cmd in commands + subroutine_types:
            docs += f"{cmd.name}:\n"
            if cmd.docstring is not None:
                docs += f"  docstring: {cmd.docstring}\n"
            if cmd.signature is not None:
                docs += f"  signature: {cmd.signature}\n"
            else:
                docs += f"  signature: {self.get_signature(cmd)}\n"
            if "arguments" in cmd.__dict__ and cmd.arguments is not None:
                docs += "  arguments:\n"
                for param, settings in cmd.arguments.items():
                    req_string = "required" if settings["required"] else "optional"
                    docs += f"    - {param} ({settings['type']}) [{req_string}]: {settings['description']}\n"
            docs += "\n"
        return docs
        