# Copyright (c) 2021 Project CHIP Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import defaultdict
from collections.abc import Iterator
import contextlib
import functools
import threading
import logging
import os
import shutil
import tarfile
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import StrEnum
from typing import Concatenate, ParamSpec, TypeVar
from collections.abc import Callable

from runner.runner import Runner


class BuildProfile(StrEnum):
    """Build profile to use for the build."""
    DEFAULT = "default"                  # Profile defined by the build system.
    DEBUG = "debug"                      # Default debug profile.
    DEBUG_OPTIMIZED = "debug-optimized"  # Debug profile with optimizations enabled.
    RELEASE = "release"                  # Release profile optimized for performance.
    RELEASE_SIZE = "release-size"        # Release profile optimized for size.


log = logging.getLogger(__name__)


@dataclass
class BuilderOptions:
    # Build profile
    build_profile: BuildProfile = BuildProfile.DEFAULT
    # Generate a link map file
    enable_link_map_file: bool = False
    # Enable flashbundle generation stage
    enable_flashbundle: bool = False
    # Allow to wrap default build command
    pw_command_launcher: str | None = None
    # Locations where files are pre-generated
    pregen_dir: str | None = None


@dataclass
class BuilderOutput:
    source: str  # Source file generated by the build
    target: str  # Target file to be copied to


class OutDirLock:
    """Lock to provide mutual exclusion for operations on output directories."""

    def __init__(self) -> None:
        self._dir_locks: dict[str, threading.RLock] = defaultdict(threading.RLock)
        self._dict_lock: threading.Lock = threading.Lock()

    @contextlib.contextmanager
    def lock_dir(self, dir_path: str | None) -> Iterator[None]:
        # No directory to lock, just yield
        if dir_path is None:
            yield
            return

        with self._dict_lock:
            dir_lock = self._dir_locks[dir_path]

        with dir_lock:
            yield


S = TypeVar('S', bound='Builder')
P = ParamSpec('P')
R = TypeVar('R')


def lock_output_dir(func: Callable[Concatenate[S, P], R | Iterator[R]]) -> Callable[Concatenate[S, P], R | Iterator[R]]:
    """Decorator to wrap build steps with output directory lock."""

    @functools.wraps(func)
    def wrapper(self: S, *args: P.args, **kwargs: P.kwargs) -> R | Iterator[R]:
        if self.output_dir_lock is None:
            return func(self, *args, **kwargs)

        output_dir_lock = self.output_dir_lock

        # Keep the lock held while an iterator result is consumed. This is needed for build_outputs and bundle_outputs which are
        # generators yielding BuilderOutput objects.
        def iterator_with_lock(result: Iterator[R]) -> Iterator[R]:
            with output_dir_lock.lock_dir(self.output_dir):
                yield from result

        with output_dir_lock.lock_dir(self.output_dir):
            return iterator_with_lock(result) if isinstance(result := func(self, *args, **kwargs), Iterator) else result

    return wrapper


class Builder(ABC):
    """Generic builder base class for CHIP.

    Provides ability to bootstrap and copy output artifacts and subclasses can
    use a generic shell runner.

    """

    def __init__(self, root: str, runner: Runner, output_dir_lock: OutDirLock | None = None):
        self.root = os.path.abspath(root)
        self._runner = runner
        self.output_dir_lock = output_dir_lock

        # Set post-init once actual build target is known
        self.identifier = None
        self.output_dir = None
        self.options = BuilderOptions()
        self.quiet = False
        self.verbose = False

    @abstractmethod
    def generate(self):
        """Generate the build files - generally the ninja/makefiles"""
        raise NotImplementedError

    @abstractmethod
    def _build(self):
        """Perform an actual build"""
        raise NotImplementedError

    def _bundle(self):
        """Perform an actual generating of flashbundle.

           May do nothing (and builder can choose not to implement this) if
           the app does not need special steps for generating flashbundle (e.g.
           on Linux platform, the output ELF files can be used directly).
        """
        pass

    @abstractmethod
    def build_outputs(self):
        """Return a list of relevant BuilderOutput objects after a build.

           May use build output data (e.g. manifests), so this should be
           invoked only after a build has succeeded.
        """
        raise NotImplementedError

    def bundle_outputs(self):
        """Return the BuilderOutput objects in flashbundle.

           Return an empty list (and builder can choose not to implement this)
           if the app does not need special files as flashbundle.

           May use data from _bundle(), so this should be invoked only after
           _bundle() has succeeded.
        """
        return []

    def outputs(self):
        outputs = list(self.build_outputs())
        if self.options.enable_flashbundle:
            outputs.extend(self.bundle_outputs())
        return outputs

    def build(self):
        self._build()
        if self.options.enable_flashbundle:
            self._bundle()

    def _Execute(self, cmdarray, title=None, dedup=False):
        self._runner.Run(cmdarray, title=title, dedup=dedup, quiet=self.quiet)

    def CompressArtifacts(self, target_file: str):
        with tarfile.open(target_file, "w:gz") as tar:
            for output in self.outputs():
                log.info('Adding %s into %s(%s)',
                         output.source, target_file, output.target)
                tar.add(output.source, output.target)

    def CopyArtifacts(self, target_dir: str):
        for output in self.outputs():
            log.info('Copying %s into %s', output.source, output.target)

            target_full_name = os.path.join(target_dir, output.target)
            target_dir_full_name = os.path.dirname(target_full_name)

            if not os.path.exists(target_dir_full_name):
                log.info('Creating subdirectory %s first',
                         target_dir_full_name)
                os.makedirs(target_dir_full_name)

            shutil.copyfile(output.source, target_full_name)
            shutil.copymode(output.source, target_full_name)