"""
Using Context in a execution graph, and specifically LoggableContext
to get rich logging and report about the execution process.

Here we simulate some kind of "Asset graph" where assets can be executed.
Each Asset type has a specific behavior when executed.
Assets are connected to other Assets input.
The name of the input they connect to is called 'input group' and can be:
- None: default connection
- build: for Builder assets
- edit: for Editor assets
- view: for Viewer assets

The system uses contexts to affect the execution of an upstream Asset from a
downstream asset.
For example:

  ┌───────┐
  │EDITOR ┼─────┐
  └───────┘     │
  ┌───────┐    ┌▼──────┐
  │BUILDER┼────► SCENE │
  └───────┘    └▲──────┘
  ┌───────┐     │
  │PRODUCT┼─────┘
  └───────┘

The SCENE needs to trigger the EDITOR when executed.
If it was never built before, it needs the BUILDER to built a file containing PRODUCT.
EDITOR and BUILDER both need to be generic so we can use it on any SCENE.

So:
- SCENE.execute() create a BuilContext(asset_to_build=self) before calling BUILDER.execute()
- BUILDER.execute() reads inputs from the context's asset_to_build field and call PRODUCT.execute()
- PRODUCT.execute() calls execute on all its inputs
- etc...
"""

from __future__ import annotations
from typing import Generator, Any

from contextlib import contextmanager

from tgzr.context import Context, LoggableContext


#
#   Asset Contexts
#


@Context.context_type
class AssetExcecutionContext(LoggableContext):
    asset: Asset

    def _log_section_summary(self) -> str:
        return f"{self.__class__.__name__} asset={self.asset.name!r} ({self.asset.__class__.__name__})"


@Context.context_type
class BuildContext(AssetExcecutionContext):
    asset: Builder
    asset_to_build: Buildable

    def _log_section_summary(self) -> str:
        summary = AssetExcecutionContext._log_section_summary(self)
        return f"{summary} asset_to_build={self.asset_to_build.name!r} ({self.asset_to_build.__class__.__name__})"


@Context.context_type
class ViewContext(AssetExcecutionContext):
    asset_to_view: Viewable

    def _log_section_summary(self) -> str:
        summary = AssetExcecutionContext._log_section_summary(self)
        return f"{summary} asset_to_view={self.asset_to_view.name!r} ({self.asset_to_view.__class__.__name__})"


@Context.context_type
class EditContext(AssetExcecutionContext):
    asset_to_edit: Viewable

    def _log_section_summary(self) -> str:
        summary = AssetExcecutionContext._log_section_summary(self)
        return f"{summary} asset_to_edit={self.asset_to_edit.name!r} ({self.asset_to_edit.__class__.__name__})"


#
#   Asset Nodes
#
class Asset:
    def __init__(self, name: str):
        self.name = name
        self._deps: dict[str | None, list[Asset]] = dict()

    def add_input(self, asset: Asset, group: str | None = None) -> None:
        try:
            assets = self._deps[group]
        except KeyError:
            assets = []
            self._deps[group] = assets
        assets.append(asset)

    def get_inputs(self, group: str | None = None) -> list[Asset]:
        return self._deps.get(group, [])

    def get_unique_input(self, group: str, required: bool = True):
        inputs = self.get_inputs(group=group)
        if required and not inputs:
            raise Exception(
                f"You need to connect something in {group!r} dependency of asset {self.name!r}!"
            )

        if len(inputs) != 1:
            raise Exception(
                f"You need one and exactly one {group!r} dependency. (got: {inputs})"
            )
        return inputs[0]

    @contextmanager
    def execution_context(self) -> Generator[AssetExcecutionContext, None, None]:
        with Context.get(AssetExcecutionContext, asset=self) as asset_context:
            yield asset_context

    def execute(self) -> None:
        raise NotImplementedError()


class AssetWithParams(Asset):
    def __init__(self, name: str):
        super().__init__(name)
        self._params: dict[str, Any] = dict()

    def set_param(self, param_name: str, param_value: Any):
        self._params[param_name] = param_value

    def set_params(self, **params):
        self._params.update(params)


class Computable(AssetWithParams):
    """
    An Asset returning a value in its 'execute()' method.
    """

    def compute(self) -> Any:
        return None

    def execute(self) -> Any:
        return self.compute()


class Buildable(AssetWithParams):
    """
    An Asset which executes its 'build' input in its context.
    """

    def scene_path(self) -> str:
        return f"/path/to/work/{self.name}/{self.name}.xstage"

    def scene_exists(self) -> bool:
        return False

    def build_if_needed(self):
        if not self.scene_exists():
            AssetExcecutionContext.current().log(
                f"Scene {self.scene_path()} is missing, we need to build it."
            )
            self.build()

    def build(self):
        builder = self.get_unique_input("build")
        with Context.get(BuildContext, asset=builder, asset_to_build=self):
            builder.execute()

    def execute(self):
        self.build_if_needed()


class Viewable(Buildable):
    """
    A Buildable Asset which executes its 'viewer' input in its context after building
    the asset if needed.
    """

    def view(self):
        with self.execution_context() as exec_ctx:
            exec_ctx.log(f"Opening viewer for {self.scene_path()}")
            viewer = self.get_unique_input("view")
            with exec_ctx.get(ViewContext, asset_to_view=self):
                viewer.execute()

    def execute(self):
        super().execute()


class Product(Viewable):
    """
    A Viewable Asset which include itself in the current scene when executed in
    a build context.
    All dependency Products are executed too.
    """

    def execute(self):
        ctx = BuildContext.current(raises=False)
        if ctx is None:
            raise Exception("Cannot execute a Product outside of a build context !")

        inputs = self.get_inputs()
        if inputs:
            with ctx.log_group(f"Executing inputs of {self.name!r}"):
                for input in inputs:
                    input.execute()

        ctx.log(f"Adding Product {self.name} to {ctx.asset_to_build.scene_path()}")


class Editable(Buildable):
    """
    A Buildable Asset which executes its 'edit' input in its context after building
    the asset if needed.
    """

    def edit(self):
        ctx = AssetExcecutionContext.current()
        editor = self.get_unique_input("edit")
        with AssetExcecutionContext.get(EditContext, asset=editor, asset_to_edit=self):
            editor.execute()

    def execute(self):
        super().execute()
        self.edit()


class Env(Computable):
    """
    A Computable Asset returning a dict[str, str] from its executes method.
    All input assets must be Env assets too, and the resulting dict will
    confain the aggregation of all input dicts.

    ! In real Asset context, the inputs order is not 100% assured !
    """

    def as_env(self, asset) -> Env | None:
        if not isinstance(asset, Env):
            return None
        return asset

    def compute(self) -> dict:
        env = dict()
        for input in self.get_inputs():
            if env_input := self.as_env(input):
                data = env_input.execute()
                env.update(data)
        for name, value in self._params.items():
            env[name] = value
        return env

    def execute(self) -> dict[str, str]:
        return {}


class Viewer(Asset):
    """
    An asset which will open the 'scene' from the current Editable asset with
    an defined viewer when executed.
    """

    def launch_viewer(self, scene: str):
        ViewContext.current().log(f"Launching viewer for {scene}")

    def execute(self):
        ctx = ViewContext.current(raises=False)
        if ctx is None:
            raise Exception("Cannot view an asset outside of a ViewContext.")
        with self.execution_context() as exec_ctx:
            exec_ctx.log(f"Viewer executing: {self.name}")
            self.launch_viewer(ctx.asset_to_view.scene_path())


class Editor(Asset):
    """
    An asset which will open the 'scene' from the current Editable asset with
    an defined DCC when executed.
    """

    def launch_editor(self, scene: str):
        EditContext.current().log(f"Launching editor on {scene}")

    def execute(self):
        ctx = EditContext.current(raises=False)
        if ctx is None:
            raise Exception("Cannot edit an asset outside of an EditContext.")
        self.launch_editor(ctx.asset_to_edit.scene_path())


class Builder(Asset):
    """
    An asset which will build the 'scene' from the current Editable asset when
    executed.
    """

    def as_buildable(self, asset) -> Buildable | None:
        if not isinstance(asset, Buildable):
            return None
        return asset

    def as_product(self, asset) -> Product | None:
        if not isinstance(asset, Product):
            return None
        return asset

    def as_env(self, asset) -> Env | None:
        if not isinstance(asset, Env):
            return None
        return asset

    def build(self):
        ctx = BuildContext.current()

        # just for testing, it's ensured in real life:
        assert isinstance(ctx.asset_to_build, Buildable)

        with ctx.log_group(
            f"Executing Products & Envs connected to {ctx.asset_to_build.name!r}"
        ):
            for input_asset in ctx.asset_to_build.get_inputs():
                if product_input := self.as_product(input_asset):
                    with ctx.log_group(f"Executing {product_input.name!r}"):
                        product_input.execute()
                elif env_input := self.as_env(input_asset):
                    with ctx.log_group(f"Executing {env_input.name!r}"):
                        env = env_input.execute()
                        ctx.log(f"Env computed: {env}")
                else:
                    ctx.log(f"Skipping non Product nor Env input: {input_asset.name}")

    def execute(self):
        build_ctx = BuildContext.current(raises=False)
        if build_ctx is None:
            raise Exception("Cannot build an asset outside of a BuildContext.")
        self.build()


class WorkScene(Editable):
    pass


def demo_asset_context():

    # Let's create an Asset graph

    alice = Product("Alice-Actor")
    bob = Product("Bob-Actor")

    kitchen = Product("Kitchen-Set")
    table = Product("Table-Prop")
    chair = Product("Chair-Prop")
    kitchen_casting = Product("Kitcher-Set-Casting")
    kitchen_casting.set_params(nb_chair=4)
    kitchen_casting.add_input(table)
    kitchen_casting.add_input(chair)
    kitchen.add_input(kitchen_casting)

    workscene = WorkScene("Anim")
    workscene.add_input(alice)
    workscene.add_input(bob)
    workscene.add_input(kitchen)

    anim_env = Env("AnimEnv")
    workscene.add_input(anim_env)

    builder = Builder("AnimBuilder")
    workscene.add_input(builder, group="build")

    editor = Editor("AnimDCC")
    workscene.add_input(editor, group="edit")

    viewer = Viewer("AnimViewer")
    workscene.add_input(viewer, group="view")

    picker_tool = Env("PickerTool")
    picker_tool.set_params(
        TOOL_PATH="/path/to/tools/picker_tool",
        PICKER_TOOL_GL="1",
    )
    anim_env.add_input(picker_tool)

    qc_tool = Env("QCTool")
    qc_tool.set_params(
        TOOL_PATH="/path/to/tools/qc_tool",
        QC_TOOL_STRICT="False",
    )
    anim_env.add_input(qc_tool)

    # Let's execute the workscene.
    # It should build itself before launchig its editor
    # Building is done by its builder, which will build all
    # products connected to the scene.
    logger = AssetExcecutionContext.create_logger(record=True)
    with Context.get(AssetExcecutionContext, logger=logger, asset=workscene):
        workscene.execute()

    # print("RECORDED:")
    # rich.print(logger.export_text())


if __name__ == "__main__":
    demo_asset_context()
