# import the contents of the Rust library into the Python extension
from .stitch_core import compress_backend,rewrite_backend
from typing import Dict, List, Any
import json

class StitchException(Exception):
    """Raised when the Stitch's Rust backend panics"""
    pass

class Abstraction:
    """
    An abstraction.

    :param name: the name of the abstraction, like "fn_0"
    :type name: str
    :param body: the body of the abstraction, like "(+ 3 (* #1 #0))". Variables are represented as "#0", "#1", etc.
    :type body: str
    :param arity: the arity of the abstraction, like 2
    :type arity: int
    """
    def __init__(self, name: str, body: str, arity: int):
        self.name = name
        self.body = body
        self.arity = arity
    def __repr__(self):
        args = ','.join([f'#{i}' for i in range(self.arity)])
        return f"{self.name}({args}) := {self.body}"

class CompressionResult:
    """
    The result of calling compress().

    :param abstractions: a list of Abstraction objects
    :type abstractions: List[Abstraction]
    :param rewritten: a list of programs, where each program has been rewritten using the abstractions
    :type rewritten: List[str]
    :param json: the raw JSON output from the Rust backend, containing lots of additional information
    :type json: Dict[str,Any]
    """
    def __init__(self, json: Dict[str,Any]):
        self.abstractions: List[Abstraction] = [Abstraction(body=abs["body"], name=abs["name"], arity=abs["arity"]) for abs in json["abstractions"]]
        self.rewritten: List[str] = json['rewritten']
        self.json = json


def from_dreamcoder(json: Dict[str,Any]):
    """
    Takes a dreamcoder-style dictionary and returns a dictionary of arguments to pass to stitch.compress()
    """

    frontiers = json["frontiers"]
    anonymous_abstractions = [production["expression"] for production in json["DSL"]["productions"] if production["expression"].startswith("#")]
    anonymous_abstractions.sort(key=len)
    anonymous_to_named = [(f"dreamcoder_abstraction_{i}", anonymous) for (i,anonymous) in enumerate(anonymous_abstractions)]

    programs = []
    tasks = []
    for i, frontier in enumerate(frontiers):
        task = frontier.get("task", str(i))
        for program in frontier["programs"]:
            program = program["program"]
            # replace #(lambda ...) with fn_2 etc. Start with highest numbered fn to avoid mangling bodies of other fns.
            for (name, anonymous) in reversed(anonymous_to_named):
                program = program.replace(anonymous, name)
            assert '#' not in program
            programs.append(program)
            tasks.append(task)

    return dict(
        programs=programs,
        tasks=tasks,
        anonymous_to_named=anonymous_to_named,
        rewritten_dreamcoder=True,
    )

def rewrite(
    programs: List[str],
    abstractions: List[Abstraction],
    **kwargs
    ) -> List[str]:
    """
    Rewrites a set of programs with a list of abstractions. Rewrites first with abstractions[0],
    then abstractions[1], etc.

    :param programs: A list of programs to rewrite.
    :type programs: List[str]
    :param abstractions: A list of Abstraction objects to rewrite with.
    :type abstractions: List[Abstraction]
    :param \**kwargs: Additional arguments to pass to the Rust backend. These are the same as :ref:`compress_kwargs` since rewriting performs a form of compression under the hood, however most are not relevant to rewriting.
    :raises StitchException: If the Rust backend panics.
    :raises TypeError: If the wrong types are provided for arguments.
    :return: A list of rewritten programs.
    :rtype: List[str]
    """

    args = " ".join([build_arg(k, v) for k, v in kwargs.items()])

    try:
        return rewrite_backend(
            programs,
            abstractions,
            args
        )
    except BaseException as e:
        if e.__class__.__name__ == "PanicException":
            raise StitchException(f"Rust backend panicked with exception: {e}")
        else:
            raise # eg TypeError from pyo3 conversion


def compress(
    programs: List[str],
    iterations: int,
    max_arity: int = 2,
    threads: int = 1,
    silent: bool = True,
    **kwargs
    ) -> CompressionResult:
    """
    Runs abstraction learning on a list of programs.

    :param programs: A list of programs to learn abstractions from.
    :type programs: List[str]
    :param iterations: The number of iterations to run abstraction learning for.
    :type iterations: int
    :param max_arity: The maximum arity of abstractions to learn.
    :type max_arity: int
    :param threads: The number of threads to use.
    :type threads: int
    :param silent: Whether to print progress to stdout.
    :type silent: bool
    :param \**kwargs: Additional arguments to pass to the Rust backend. See :ref:`compress_kwargs` for a full listing.
    :raises StitchException: If the Rust backend panics.
    :raises TypeError: If the wrong types are provided for arguments.
    :return: A CompressionResult object containing the learned abstractions, rewritten programs, and other details from the run.
    :rtype: CompressionResult
    """

    tasks = kwargs.pop("tasks", None)
    anonymous_to_named = kwargs.pop("anonymous_to_named", None)

    kwargs.update(dict(
        iterations=iterations,
        max_arity=max_arity,
        threads=threads,
        silent=silent
    ))

    args = " ".join([build_arg(k, v) for k, v in kwargs.items()])

    try:
        res = compress_backend(
            programs,
            tasks,
            anonymous_to_named,
            args)
    except BaseException as e:
        if e.__class__.__name__ == "PanicException":
            raise StitchException(f"Rust backend panicked with exception: {e}")
        else:
            raise # eg TypeError from pyo3 conversion
    
    res = json.loads(res)
    
    return CompressionResult(res)
    

def build_arg(name: str, val) -> str:
    """
    Builds command line argument version of a Python argument, so for example:
    - build_arg("max_arity",3) -> "--max-arity=3"
    - build_arg("silent",True) -> "--silent"
    - build_arg("silent",False) -> ""
    """
    assert isinstance(val, (str, int, bool)), f"unexpected type for argument `{name}`: {type(val)}"
    res = "--" + name.replace("_", "-")
    if isinstance(val,bool):
        if val:
            return res # eg "--silent"
        else:
            return ""
    return f"{res}={val}"