"""
Copyright 2024 HaiyangLi

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 abc import ABC
from functools import wraps
from typing import Callable
import asyncio
import inspect
from lionagi import logging as _logging
from lionagi.libs.ln_func_call import pcall
from lionagi.core.work.work_function import WorkFunction
from lionagi.core.work.work import Work
from lionagi.core.report.form import Form
from lionagi.core.collections.abc import get_lion_id


class Worker(ABC):
    """
    This class represents a worker that handles multiple work functions.

    Attributes:
        name (str): The name of the worker.
        work_functions (dict[str, WorkFunction]): Dictionary mapping assignments to WorkFunction objects.
        forms (dict[str, Form]): Dictionary mapping form identifier to Form objects.
        default_form (str|None): The default form to be used by the worker.
    """

    name: str = "Worker"
    work_functions: dict[str, WorkFunction] = {}

    def __init__(self, forms=None, default_form=None) -> None:
        # self.stopped = False
        self.forms: dict[str, Form] = forms or {}
        self.default_form = default_form

    async def stop(self):
        """
        Stops the worker and all associated work functions.
        """
        # self.stopped = True
        _logging.info(f"Stopping worker {self.name}")
        non_stopped_ = []

        for func in self.work_functions.values():
            worklog = func.worklog
            await worklog.stop()
            if not worklog.stopped:
                non_stopped_.append(func.name)

        if len(non_stopped_) > 0:
            _logging.error(f"Could not stop worklogs: {non_stopped_}")
        _logging.info(f"Stopped worker {self.name}")

    async def is_progressable(self):
        """
        Checks if any work function is progressable and the worker is not stopped.

        Returns:
            bool: True if any work function is progressable and the worker is not stopped, else False.
        """

        return (
            any([await i.is_progressable() for i in self.work_functions.values()])
            and not self.stopped
        )

    async def change_default_form(self, form_key):
        """
        Changes the default form to the specified form key.

        Args:
            form_key (str): The key of the form to set as the default.

        Raises:
            ValueError: If the form key does not exist in the forms dictionary.

        """
        if form_key not in self.forms.keys():
            raise ValueError(f"Unable to change default form. Key {form_key} does not exist.")
        self.default_form = self.forms[form_key]

    # async def process(self, refresh_time=1):
    #     """
    #     Processes all work functions periodically.
    #
    #     Args:
    #         refresh_time (int): Time interval between each process cycle.
    #     """
    #     while await self.is_progressable():
    #         await pcall([i.process(refresh_time) for i in self.work_functions.values()])
    #         await asyncio.sleep(refresh_time)

    # TODO: Implement process method

    # async def process(self, refresh_time=1):
    #     while not self.stopped:
    #         tasks = [
    #             asyncio.create_task(func.process(refresh_time=refresh_time))
    #             for func in self.work_functions.values()
    #         ]
    #         await asyncio.wait(tasks)
    #         await asyncio.sleep(refresh_time)

    async def _wrapper(
        self,
        *args,
        func=None,
        assignment=None,
        form_param_key=None,
        capacity=None,
        retry_kwargs=None,
        guidance=None,
        **kwargs,
    ):
        """
        Internal wrapper to handle work function execution.

        Args:
            func (Callable): The function to be executed.
            assignment (str): The assignment description.
            form_param_key (str): The key to identify the form parameter in
                the function's signature. This parameter is used to locate and fill
                the appropriate form according to the assignment. Raises an error
                if the form parameter key is not found in the function's signature.
            capacity (int): Capacity for the work log.
            retry_kwargs (dict): Retry arguments for the function.
            guidance (str): Guidance or documentation for the function.
        """
        if getattr(self, "work_functions", None) is None:
            self.work_functions = {}

        if func.__name__ not in self.work_functions:
            self.work_functions[func.__name__] = WorkFunction(
                assignment=assignment,
                function=func,
                retry_kwargs=retry_kwargs or {},
                guidance=guidance,
                capacity=capacity,
            )

        work_func: WorkFunction = self.work_functions[func.__name__]

        # locate form that should be filled according to the assignment
        if form_param_key:
            func_signature = inspect.signature(func)
            if form_param_key not in func_signature.parameters:
                raise KeyError(f"Failed to locate form. \"{form_param_key}\" is not defined in the function.")
            if "self" in func_signature.parameters:
                bound_args = func_signature.bind(None, *args, **kwargs)
            else:
                bound_args = func_signature.bind(*args, **kwargs)
            bound_args.apply_defaults()
            arguments = bound_args.arguments

            form_key = arguments.get(form_param_key)
            try:
                form_key = get_lion_id(form_key)
            except:
                pass
            form = self.forms.get(form_key) or self.default_form

            if form:
                subform = form.__class__(assignment=work_func.assignment, task=work_func.guidance)
                for k in subform.input_fields:
                    v = getattr(form, k, None)
                    setattr(subform, k, v)
                subform.origin = form
                kwargs = {"form": subform} | kwargs
            else:
                raise ValueError(f"Cannot locate form in Worker's forms and default_form is not available.")

        task = work_func.perform(self, *args, **kwargs)
        work = Work(async_task=task, async_task_name=work_func.name)
        await work_func.worklog.append(work)
        return work


def work(
    assignment=None,
    form_param_key=None,
    capacity=10,
    guidance=None,
    retry_kwargs=None,
    timeout=10,
):
    """
    Decorator to mark a method as a work function.

    Args:
        assignment (str): The assignment description of the work function.
        form_param_key (str): The key to identify the form parameter in
                the function's signature. This parameter is used to locate and fill
                the appropriate form according to the assignment. Raises an error
                if the form parameter key is not found in the function's signature.
        capacity (int): Capacity for the work log.
        guidance (str): Guidance or documentation for the work function.
        retry_kwargs (dict): Retry arguments for the work function.
        timeout (int): Timeout for the work function.
    """

    def decorator(func):
        @wraps(func)
        async def wrapper(
            self: Worker,
            *args,
            func=func,
            assignment=assignment,
            form_param_key=form_param_key,
            capacity=capacity,
            retry_kwargs=retry_kwargs,
            guidance=guidance,
            **kwargs,
        ):
            retry_kwargs = retry_kwargs or {}
            retry_kwargs["timeout"] = retry_kwargs.get("timeout", timeout)
            return await self._wrapper(
                *args,
                func=func,
                assignment=assignment,
                form_param_key=form_param_key,
                capacity=capacity,
                retry_kwargs=retry_kwargs,
                guidance=guidance,
                **kwargs,
            )

        return wrapper

    return decorator


# # Example
# from lionagi import Session
# from lionagi.experimental.work.work_function import work


# class MyWorker(Worker):

#     @work(assignment="instruction, context -> response")
#     async def chat(instruction=None, context=None):
#         session = Session()
#         return await session.chat(instruction=instruction, context=context)


# await a.chat(instruction="Hello", context={})
