"""
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 enum import Enum
from lionagi.libs.ln_convert import to_str, to_dict
from lionagi.core.collections.abc import Field
from .template.base import BaseUnitForm


class UnitForm(BaseUnitForm):
    """
    Form for managing unit directives and outputs.

    Attributes:
        actions (dict|None): Actions to take based on the context and instruction.
        action_required (bool|None): Set to True if actions are provided, else False.
        answer (str|None): Answer to the questions asked.
        extension_required (bool|None): Set to True if more steps are needed.
        prediction (str|None): Likely prediction based on context and instruction.
        plan (dict|str|None): Step-by-step plan.
        next_steps (dict|str|None): Ideas on next actions to take.
        score (float|None): Numeric performance score.
        reflection (str|None): Self-reflection on reasoning and performance.
        selection (Enum|str|list|None): A single item from the choices.
        tool_schema (list|dict|None): The list of tools available for use.
        assignment (str): Default assignment task description.
    """

    actions: dict | None = Field(
        None,
        description=(
            "Actions to take based on the context and instruction. "
            "Format: {action_n: {function: ..., arguments: {param1: ..., "
            "param2: ...}}}. Leave blank if no actions are needed. "
            "Must use provided functions and parameters, DO NOT MAKE UP NAMES!!! "
            "Flag `action_required` as True if filled. When providing parameters, "
            "you must follow the provided type and format. If the parameter is a "
            "number, you should provide a number like 1, 23, or 1.1 if float is "
            "allowed."
        ),
        examples=["{action_1: {function: 'add', arguments: {num1: 1, num2: 2}}}"],
    )

    action_required: bool | None = Field(
        None,
        description="Set to True if you provide actions, or provide actions if True.",
        examples=[True, False],
    )

    answer: str | None = Field(
        None,
        description=(
            "Adhere to the prompt and all user instructions. Provide the answer "
            "for the task. if actions are required at this step, set "
            "`action_required` to True and write only `PLEASE_ACTION` to the answer_field."
            "Additionally, if extensions are allowed and needed at this step to provide a"
            " high-quality, accurate answer, set extension_required to True and"
            "you will have another chance to provide the answer after the actions are done.",
        ),
    )

    extension_required: bool | None = Field(
        None,
        description=(
            "Set to True if more steps are needed to provide an accurate answer. "
            "If True, additional rounds are allowed."
        ),
        examples=[True, False],
    )

    prediction: None | str = Field(
        None,
        description="Provide the likely prediction based on context and instruction.",
    )

    plan: dict | str | None = Field(
        None,
        description=(
            "Provide a step-by-step plan. Format: {step_n: {plan: ..., reason: ...}}. "
            "Achieve the final answer at the last step. Set `extend_required` to True "
            "if plan requires more steps."
        ),
        examples=["{step_1: {plan: '...', reason: '...'}}"],
    )

    next_steps: dict | str | None = Field(
        None,
        description=(
            "Brainstorm ideas on next actions to take. Format: {next_step_n: {plan: "
            "..., reason: ...}}. Next_step is about anticipating future actions, "
            "but it does not have to be in a sequential format. Set `extend_required` "
            "to True if more steps are needed."
        ),
        examples=["{next_step_1: {plan: '...', reason: '...'}}"],
    )

    score: float | None = Field(
        None,
        description=(
            "A numeric score. Higher is better. If not otherwise instructed, "
            "fill this field with your own performance rating. Try hard and be "
            "self-critical."
        ),
        examples=[0.2, 5, 2.7],
    )

    reflection: str | None = Field(
        None,
        description=(
            "Reason your own reasoning. Create specific items on how/where/what "
            "you could improve to better achieve the task, or can the problem be "
            "solved in a different and better way. If you can think of a better "
            "solution, please provide it and fill the necessary fields in "
            "`action_required`, `extension_required`, `next_steps` if appropriate."
        ),
    )

    selection: Enum | str | list | None = Field(
        None, description="a single item from the choices."
    )

    tool_schema: list | dict | None = Field(
        None, description="The list of tools available for using."
    )

    assignment: str = Field("task -> answer")

    is_extension: bool = Field(False)

    def __init__(
        self,
        instruction=None,
        context=None,
        reason: bool = True,
        predict: bool = False,
        score=True,
        select=None,
        plan=None,
        brainstorm=None,
        reflect=None,
        tool_schema=None,
        allow_action: bool = False,
        allow_extension: bool = False,
        max_extension: int = None,
        confidence=None,
        score_num_digits=None,
        score_range=None,
        select_choices=None,
        plan_num_step=None,
        predict_num_sentences=None,
        **kwargs,
    ):
        """
        Initialize the UnitForm with various parameters and settings.

        Args:
            instruction (str|None): Additional instruction.
            context (str|None): Additional context.
            reason (bool): Flag to include reasoning.
            predict (bool): Flag to include prediction.
            score (bool): Flag to include score.
            select (Enum|str|list|None): Selection choices.
            plan (dict|str|None): Step-by-step plan.
            brainstorm (bool): Flag to include brainstorming next steps.
            reflect (bool): Flag to include self-reflection.
            tool_schema (list|dict|None): Schema of available tools.
            allow_action (bool): Allow actions to be added.
            allow_extension (bool): Allow extension for more steps.
            max_extension (int|None): Maximum number of extensions allowed.
            confidence (bool|None): Flag to include confidence score.
            score_num_digits (int|None): Number of decimal places for the score.
            score_range (list|None): Range for the score.
            select_choices (list|None): Choices for selection.
            plan_num_step (int|None): Number of steps in the plan.
            predict_num_sentences (int|None): Number of sentences for prediction.
            **kwargs: Additional keyword arguments.
        """
        super().__init__(**kwargs)

        self.task = (
            f"Follow the prompt and provide the necessary output.\n"
            f"- Additional instruction: {to_str(instruction or 'N/A')}\n"
            f"- Additional context: {to_str(context or 'N/A')}\n"
        )

        if reason:
            self.append_to_request("reason")

        if allow_action:
            self.append_to_request("actions, action_required, reason")
            self.task += "- Reason and prepare actions with GIVEN TOOLS ONLY.\n"

        if allow_extension:
            self.append_to_request("extension_required")
            self.task += (
                f"- Allow auto-extension up to another {max_extension} rounds.\n"
            )

        if tool_schema:
            self.append_to_input("tool_schema")
            self.tool_schema = tool_schema

        if brainstorm:
            self.append_to_request("next_steps, extension_required")
            self.task += "- Explore ideas on next actions to take.\n"

        if plan:
            plan_num_step = plan_num_step or 3
            max_extension = max_extension or plan_num_step
            allow_extension = True
            self.append_to_request("plan, extension_required")
            self.task += (
                f"- Generate a {plan_num_step}-step plan based on the context.\n"
            )

        if predict:
            self.append_to_request("prediction")
            self.task += (
                f"- Predict the next {predict_num_sentences or 1} sentence(s).\n"
            )

        if select:
            self.append_to_request("selection")
            self.task += (
                f"- Select 1 item from the provided choices: {select_choices}.\n"
            )

        if confidence:
            self.append_to_request("confidence_score")

        if score:
            self.append_to_request("score")

            score_range = score_range or [0, 10]
            score_num_digits = score_num_digits or 0

            self.validation_kwargs["score"] = {
                "upper_bound": score_range[1],
                "lower_bound": score_range[0],
                "num_type": int if score_num_digits == 0 else float,
                "precision": score_num_digits if score_num_digits != 0 else None,
            }

            self.task += (
                f"- Give a numeric score in [{score_range[0]}, {score_range[1]}] "
                f"and precision of {score_num_digits or 0} decimal places.\n"
            )

        if reflect:
            self.append_to_request("reflection")

    def display(self):
        """
        Display the current form fields and values in a user-friendly format.
        """
        fields = self.work_fields.copy()

        if "task" in fields and len(str(fields["task"])) > 2000:
            fields["task"] = fields["task"][:2000] + "..."

        if "tool_schema" in fields:
            tools = to_dict(fields["tool_schema"])["tools"]
            fields["available_tools"] = [i["function"]["name"] for i in tools]
            fields.pop("tool_schema")

        if "actions" in fields:
            a = ""
            idx = 0
            for _, v in fields["actions"].items():
                a += (
                    f"\n \n{idx+1}. **{v['function']}**"
                    f"({', '.join([f'{k}: {v}' for k, v in v['arguments'].items()])}), "
                )
                idx += 1
            fields["actions"] = a[:-2]

        if "action_response" in fields:
            a = ""
            idx = 0
            for _, v in fields["action_response"].items():
                a += (
                    f"\n \n{idx+1}. **{v['function']}**"
                    f"({', '.join([f'{k}: {v}' for k, v in v['arguments'].items()])})"
                )
                if len(str(v["output"])) > 30:
                    a += f" \n \n {v['output']}, "
                else:
                    a += f" = {v['output']}, "
                idx += 1
            fields["action_response"] = a[:-2]

        super().display(fields=fields)
