import re
from datetime import datetime, timezone
from functools import wraps
from importlib.resources import files
from uuid import uuid4

import yaml
from relai_visualization.models import RELAIData
from relai_visualization.models.application_result import ApplicationResult

from .relai_client import RELAIClient
from .utils import generate_random_phrase


class RELAIApplication:
    """Base class for all RELAI applications. All RELAI applications must inherit from this class.
    Expects a `config.yaml` file in the application package to provide information about the application.
    See https://github.com/relai-apps/template for more details.

    Args:
        project_name (str): Name of the project.
        run_name (str): Name of the run. If not provided, a random name is generated.
        schema_version (str): Version of the schema. Default is "v1.0".
        **metadata: Additional metadata for the application.
    """

    def __init__(
        self,
        project_name: str,
        schema_version: str = "v1.0",
        **metadata: dict,
    ) -> None:
        if not re.match("^[a-z0-9-_]+$", project_name):
            raise ValueError("Project name must only contain lowercase characters, hyphens, or underscores")
        self._relai_client = RELAIClient()

        # Read `config.yaml`
        package_name = self.__class__.__module__.split(".")[0]
        config_path = files(package_name) / "config.yaml"
        with open(config_path) as f:
            config = yaml.safe_load(f)

        self._application = config["application"]
        self._project_name = project_name
        self._run_id = None

        metadata["application"] = self._application  # add application name to metadata

        self._partial_application_result = {
            "schema_version": schema_version,
            "application": self._application,
            "application_version": config["application_version"],
            "display_name": config["display_name"],
            "application_info": config["application_info"],
            "project_name": project_name,
            "metadata": metadata,
        }

        self._components = {}

    def register_run(self, run_name: str = None) -> None:
        """Registers a new application run and retrieves the run id.
        All RELAI applications must call this method at the beginning of an analysis.

        Args:
            run_name (str, optional): Name of the run. If not provided, a random name is generated.
        """
        timestamp = datetime.now(timezone.utc)
        if not run_name:
            run_name = generate_random_phrase()  # use a random run name if not provided

        self._partial_application_result["started_at"] = timestamp.isoformat()

        self._run_id = self._relai_client.get_run_id(self._project_name, run_name)

    @property
    def run_id(self) -> str:
        """Returns the run id for the current application run.

        Returns:
            str: The run id.
        """
        if self._run_id is None:
            raise RuntimeError("Run ID is not set. Please call `register_run` first.")
        return self._run_id

    @property
    def components(self) -> dict[str, RELAIData]:
        """Returns the components for the current application run as a dictionary mapping component ids
        to RELAIData objects.
        """
        return self._components

    def add_component(self, component: RELAIData) -> str:
        """Adds a component to the current application run and returns the component id.

        Args:
            component (RELAIData): Component to add.

        Returns:
            str: Component id.
        """
        component_id = f"{component.data['componentName']}-{uuid4().hex[:8]}"
        self._components[component_id] = component
        return component_id

    def remove_component(self, component_id: str) -> None:
        """Removes a component from the current application run.

        Args:
            component_id (str): Component id to remove.
        """
        self._components.pop(component_id, None)

    def log(
        self,
        num_samples_processed: int,
        completed: bool = False,
    ) -> None:
        """Logs RELAI credits and data for visualization on the RELAI dashboard.

        Args:
            num_samples_processed (int): samples processed in this analysis.
            completed (bool, optional): Whether the run has completed.
        """
        self._partial_application_result["run_id"] = self.run_id
        if completed:
            # Add the end time to the application result
            self._partial_application_result["ended_at"] = datetime.now(timezone.utc).isoformat()

        application_result = ApplicationResult(**self._partial_application_result, components=self.components)
        self._relai_client.log_result(application_result, self.run_id, num_samples_processed)

    @classmethod
    def relai_analysis(cls, func):
        """Decorator to register a method as a RELAI analysis function. It adds an additional argument `run_name`
        and registers a run with the given `run_name` before calling the decorated analysis function. This decorator
        must be used on methods that perform an analysis and call the `log` method.

        Example:

        ```python
        class SampleRELAIApplication(RELAIApplication):
            @RELAIApplication.relai_analysis
            def sample_analysis(self):
                pass
        ```

        Args:
            func (callable): The function to be decorated.
        """

        @wraps(func)
        def wrapper(self: "RELAIApplication", *args, run_name: str = None, **kwargs):
            self.register_run(run_name=run_name)
            self._relai_client.log_metadata(self._partial_application_result, self._run_id)
            result = func(self, *args, **kwargs)
            return result

        return wrapper

    @classmethod
    def consumes_feedback(cls, required: bool = True):
        """Decorator to register a method as a RELAI analysis function that consumes user feedback. It adds an
        additional argument `consume_feedback_from` that specifies the run id from which the feedback should be
        consumed. This feedback will be passed to the decorated analysis function through the argument `user_feedback`.
        The `required` argument specifies whether the `consume_feedback_from` argument is required.

        Example:

            ```python
            class SampleRELAIApplication(RELAIApplication):
                @RELAIApplication.consumes_feedback(required=True)
                @RELAIApplication.relai_analysis
                def sample_analysis(self, user_feedback):
                    pass
            ```
        """

        def decorator(func):

            @wraps(func)
            def wrapper(self: "RELAIApplication", *args, consume_feedback_from: str = None, **kwargs):
                if required and consume_feedback_from is None:
                    raise ValueError("consume_feedback_from is a required argument")
                user_feedback = None
                if consume_feedback_from:
                    # TODO: Maybe filter for feedback specifically. For now, just pass the entire result.
                    # user_feedback = ApplicationResult.user_feedback(self._relai_client.get_result(consume_feedback_from))
                    user_feedback = self._relai_client.get_result(consume_feedback_from)
                result = func(self, *args, user_feedback=user_feedback, **kwargs)
                return result

            return wrapper

        return decorator
