# -*- coding: utf-8 -*-
"""
@author:XuMing(xuming624@qq.com)
@description:
part of the code from https://github.com/phidatahq/phidata
"""
from __future__ import annotations

from typing import List, Any, Optional, Dict, Iterator, Union
from uuid import uuid4

from pydantic import BaseModel, ConfigDict, field_validator, Field

from agentica.llm.base import LLM
from agentica.message import get_text_from_message
from agentica.task import Task
from agentica.utils.log import logger, set_log_level_to_debug, print_llm_stream


class Workflow(BaseModel):
    # -*- settings
    # LLM to use for this Workflow
    llm: Optional[LLM] = None
    # Workflow name
    name: Optional[str] = None

    # -*- Run settings
    # Run UUID (autogenerated if not set)
    run_id: Optional[str] = Field(None, validate_default=True)
    # Metadata associated with this run
    run_data: Optional[Dict[str, Any]] = None

    # -*- User settings
    # ID of the user running this workflow
    user_id: Optional[str] = None
    # Metadata associated the user running this workflow
    user_data: Optional[Dict[str, Any]] = None

    # -*- Tasks in this workflow (required)
    tasks: List[Task]
    # Metadata associated with the assistant tasks
    task_data: Optional[Dict[str, Any]] = None

    # -*- Workflow Output
    # Final output of this Workflow
    output: Optional[Any] = None

    # debug_mode=True enables debug logs
    debug_mode: bool = False

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @field_validator("debug_mode", mode="before")
    def set_log_level(cls, v: bool) -> bool:
        if v:
            set_log_level_to_debug()
            logger.debug("Debug logs enabled")
        return v

    @field_validator("run_id", mode="before")
    def set_run_id(cls, v: Optional[str]) -> str:
        return v if v is not None else str(uuid4())

    def _run(
            self,
            message: Optional[Union[List, Dict, str]] = None,
            *,
            stream: bool = True,
            print_output: bool = False,
            **kwargs: Any,
    ) -> Iterator[str]:
        logger.debug(f"*********** Workflow Run Start: {self.run_id} ***********")

        # List of tasks that have been run
        executed_tasks: List[Task] = []
        workflow_output: List[str] = []

        # -*- Generate response by running tasks
        for idx, task in enumerate(self.tasks, start=1):
            logger.debug(f"*********** Task {idx} Start ***********")

            # -*- Prepare input message for the current_task
            task_input: List[str] = []
            if message is not None:
                task_input.append(get_text_from_message(message))

            if len(executed_tasks) > 0:
                previous_task_outputs = []
                for previous_task_idx, previous_task in enumerate(executed_tasks, start=1):
                    previous_task_output = previous_task.get_task_output_as_str()
                    if previous_task_output is not None:
                        previous_task_outputs.append(
                            (previous_task_idx, previous_task.description, previous_task_output)
                        )

                if len(previous_task_outputs) > 0:
                    task_input.append("\nHere are previous tasks and and their results:\n---")
                    for previous_task_idx, previous_task_description, previous_task_output in previous_task_outputs:
                        task_input.append(f"Task {previous_task_idx}: {previous_task_description}")
                        task_input.append(previous_task_output)
                    task_input.append("---")

            # -*- Run Task
            task_output = ""
            input_for_current_task = "\n".join(task_input)
            if stream and task.streamable:
                for chunk in task.run(message=input_for_current_task, stream=True, **kwargs):
                    task_output += chunk if isinstance(chunk, str) else ""
                    chunk_str = chunk if isinstance(chunk, str) else ""
                    if print_output:
                        print_llm_stream(chunk_str)
                    yield chunk_str
            else:
                task_output = task.run(message=input_for_current_task, stream=False, **kwargs)  # type: ignore

            executed_tasks.append(task)
            workflow_output.append(task_output)
            logger.debug(f"*********** Task {idx} End ***********")
            if not stream:
                if print_output:
                    print(task_output)
                yield task_output
        logger.debug(f"*********** Workflow Run End: {self.run_id} ***********")

    def run(
            self,
            message: Optional[Union[List, Dict, str]] = None,
            *,
            stream: bool = True,
            print_output: bool = False,
            **kwargs: Any,
    ) -> Union[Iterator[str], str]:
        if stream:
            resp = self._run(message=message, stream=True, print_output=print_output, **kwargs)
            return resp
        else:
            resp = "".join(self._run(message=message, stream=False, print_output=print_output, **kwargs))
            return resp

    def cli(
            self,
            user: str = "User",
            emoji: str = ":sunglasses:",
            stream: bool = True,
            print_output: bool = True,
            exit_on: Optional[List[str]] = None,
            **kwargs: Any,
    ) -> None:
        from rich.prompt import Prompt
        _exit_on = exit_on or ["exit", "quit", "bye"]
        logger.debug(f"Enable cli, exit with {_exit_on[0]}")
        while True:
            message = Prompt.ask(f"[bold] {emoji} {user} [/bold]")
            if message in _exit_on:
                break

            r = self.run(message=message, stream=stream, print_output=print_output, **kwargs)
            if stream:
                print("".join(r))
            else:
                print(r)
