from typing import List, Optional
from metaflow.exception import MetaflowException
from metaflow.decorators import StepDecorator, FlowDecorator
from metaflow import current
from metaflow.metaflow_config import KUBERNETES_CONTAINER_IMAGE
from .deployer import AppDeployer, DeployedApp, PackagedCode
from .perimeters import PerimeterExtractor
import os


class AppDeployFlowDecorator(FlowDecorator):
    """
    Deploy and manage Outerbounds Apps from within a Metaflow flow.

    This decorator enables deploying apps during flow execution, automatically
    tagging them with Metaflow metadata (flow name, run ID, step, etc.) for
    easy tracking and management.

    When applied to a flow, it exposes `current.apps` which provides:
    - Access to the flow's code package and image
    - Ability to list apps deployed in the current run
    - Automatic tagging with Metaflow context

    Examples
    --------

    ```python
    from metaflow import FlowSpec, step, current, app_deploy
    from metaflow.apps import AppDeployer

    @app_deploy
    class MyFlow(FlowSpec):

        @step
        def start(self):
            # Deploy an app using the flow's code package
            deployer = AppDeployer(
                name="my-service",
                port=8000,
                image=current.apps.current_image,
                code_package=current.apps.metaflow_code_package,
                commands=["python server.py"],
            )
            self.app = deployer.deploy()
            self.next(self.end)

        @step
        def end(self):
            # List all apps deployed in this run
            apps = current.apps.list()
            print(f"Deployed {len(apps)} app(s)")
    ```


    MF Add To Current
    -----------------
    apps -> metaflow_extensions.outerbounds.plugins.apps.core.app_deploy_decorator.FlowAppManager

        @@ Returns
        ----------
        FlowAppManager
    """

    name = "app_deploy"
    defaults = {}

    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
        self._attributes_with_user_values = (
            set(attributes.keys()) if attributes is not None else set()
        )

        super().__init__(attributes, statically_defined, inserted_by)

    def flow_init(
        self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
    ):

        from metaflow import decorators

        decorators._attach_decorators(flow, ["app_deploy_internal"])
        decorators._process_late_attached_decorator(
            ["app_deploy_internal"],
            flow,
            graph,
            environment,
            flow_datastore,
            logger,
        )


class AppDeployInternalDecorator(StepDecorator):

    name = "app_deploy_internal"
    defaults = {}

    packaged_code = None

    package_url = None
    package_sha = None

    MAX_ENTROPY = 6
    MAX_NAME_LENGTH = 150

    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
        self.logger = logger
        self.environment = environment
        self.step = step
        self.flow_datastore = flow_datastore

    def _resolve_package_url_and_sha(self):
        return os.environ.get("METAFLOW_CODE_URL", self.package_url), os.environ.get(
            "METAFLOW_CODE_SHA", self.package_sha
        )

    def _extract_project_info(self):
        project = current.get("project_name")
        branch = current.get("branch_name")
        is_production = current.get("is_production")
        return project, branch, is_production

    def task_pre_step(
        self,
        step_name,
        task_datastore,
        metadata,
        run_id,
        task_id,
        flow,
        graph,
        retry_count,
        max_user_code_retries,
        ubf_context,
        inputs,
    ):
        perimeter, api_server = PerimeterExtractor.during_programmatic_access()
        package_url, package_sha = self._resolve_package_url_and_sha()
        if package_url is None or package_sha is None:
            # TODO: Think through if we need this abstraction
            raise MetaflowException(
                "METAFLOW_CODE_URL or METAFLOW_CODE_SHA is not set. "
                "Please set METAFLOW_CODE_URL and METAFLOW_CODE_SHA in your environment."
            )
        self.packaged_code = PackagedCode(package_url, package_sha)
        image = os.environ.get("FASTBAKERY_IMAGE", None)

        project, branch, is_production = self._extract_project_info()

        project_info = {}
        if project is not None:
            project_info["metaflow/project"] = project
            project_info["metaflow/branch"] = branch
            project_info["metaflow/is_production"] = is_production

        default_tags = {
            "metaflow/flow_name": flow.name,
            "metaflow/step_name": step_name,
            "metaflow/run_id": run_id,
            "metaflow/task_id": task_id,
            "metaflow/retry_count": retry_count,
            "metaflow/pathspec": current.pathspec,
            **project_info,
        }

        # Set state on AppDeployer using the new API
        AppDeployer._set_state("perimeter", perimeter)
        AppDeployer._set_state("api_url", api_server)
        AppDeployer._set_state(
            "default_tags", [{k: str(v)} for k, v in default_tags.items()]
        )

        current._update_env(
            {
                "apps": FlowAppManager(
                    flow.name,
                    run_id,
                    self.packaged_code,
                    image,
                )
            }
        )

    def task_post_step(
        self, step_name, flow, graph, retry_count, max_user_code_retries
    ):
        pass

    def runtime_init(self, flow, graph, package, run_id):
        # Set some more internal state.
        self.flow = flow
        self.graph = graph
        self.package = package
        self.run_id = run_id

    def runtime_task_created(
        self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
    ):
        # To execute the Kubernetes job, the job container needs to have
        # access to the code package. We store the package in the datastore
        # which the pod is able to download as part of it's entrypoint.
        if not is_cloned:
            self._save_package_once(self.flow_datastore, self.package)

    @classmethod
    def _save_package_once(cls, flow_datastore, package):
        if cls.packaged_code is None:
            cls.package_url, cls.package_sha = flow_datastore.save_data(
                [package.blob], len_hint=1
            )[0]
            cls.packaged_code = PackagedCode(cls.package_url, cls.package_sha)
            os.environ["METAFLOW_CODE_URL"] = cls.package_url
            os.environ["METAFLOW_CODE_SHA"] = cls.package_sha


class FlowAppManager:
    """
    Manager for apps deployed within a Metaflow flow execution.

    Accessible via `current.apps` when using the `@app_deploy` decorator.
    Provides access to the flow's code package, container image, and
    methods to list apps deployed in the current run.

    Attributes
    ----------
    metaflow_code_package : PackagedCode
        The code package for the current flow, ready to use with AppDeployer.
    current_image : str
        The container image used by the current task (from fast_bakery or similar).
    default_image : str
        The default Kubernetes container image from Metaflow config.

    Examples
    --------

    ```python
    # python myflow.py --environment=fast-bakery run --with kubernetes
    from metaflow.apps import AppDeployer

    @pypi(packages={"flask": ">=2.0", "requests": ">=2.28"})
    @step
    def deploy_step(self):
        image = current.apps.current_image
        if image is None:
            image = current.apps.default_image
        # Use the flow's code package directly
        deployer = AppDeployer(
            name="my-app",
            port=8000,
            image=image,
            code_package=current.apps.metaflow_code_package,
            commands=["python app.py"],
        )
        deployed = deployer.deploy()

        # List apps from this run
        apps = current.apps.list()
    ```
    """

    def __init__(
        self,
        flow_name: str,
        run_id: str,
        package: PackagedCode,
        image: Optional[str] = None,
    ) -> None:
        self._package = package
        self._image = image
        self._runid = run_id
        self._flowname = flow_name

    @property
    def metaflow_code_package(self) -> PackagedCode:
        """
        The flow's code package for use with AppDeployer.

        Returns
        -------
        PackagedCode
            A namedtuple with `url` and `key` fields pointing to the
            packaged metaflow's code package stored in the datastore.
        """
        return self._package

    @property
    def current_image(self) -> str:
        """
        The container image for the current task.

        Returns
        -------
        str
            Image URI from fast bakery image or None if not set.
        """
        return self._image

    @property
    def default_image(self) -> str:
        """
        The default Kubernetes container image from Metaflow config.

        Returns
        -------
        str
            The KUBERNETES_CONTAINER_IMAGE from Metaflow configuration.
        """
        return KUBERNETES_CONTAINER_IMAGE

    def list(self) -> List["DeployedApp"]:
        """
        List apps deployed in the current Metaflow run.

        Returns apps tagged with this flow's name and run ID.

        Returns
        -------
        List[DeployedApp]
            Apps deployed during this flow execution.

        Examples
        --------

        ```python
        apps = current.apps.list()
        for app in apps:
            print(f"{app.name}: {app.public_url}")
        ```
        """
        return AppDeployer.list_deployments(
            tags=[
                {"metaflow/run_id": self._runid},
                {"metaflow/flow_name": self._flowname},
            ]
        )
