from abc import abstractmethod
from dataclasses import dataclass
from typing import List, Optional

from llama_index import Document, ServiceContext

from llama_agi.default_task_prompts import (
    DEFAULT_TASK_PRIORITIZE_TMPL,
    DEFAULT_REFINE_TASK_PRIORITIZE_TMPL,
    DEFAULT_TASK_CREATE_TMPL,
    DEFAULT_REFINE_TASK_CREATE_TMPL,
)


@dataclass
class LlamaTaskPrompts:
    task_create_qa_template: str = DEFAULT_TASK_CREATE_TMPL
    task_create_refine_template: str = DEFAULT_REFINE_TASK_CREATE_TMPL
    task_prioritize_qa_template: str = DEFAULT_TASK_PRIORITIZE_TMPL
    task_prioritize_refine_template: str = DEFAULT_REFINE_TASK_PRIORITIZE_TMPL


class BaseTaskManager:
    """Base Task Manager

    Args:
        tasks (List[str]): The initial list of tasks to complete.
        prompts: (LlamaTaskPrompts): The prompts to control the task creation
        and prioritization.
        tasK_service_context (ServiceContext): The LlamaIndex service context to use
        for task creation and prioritization.

    """

    def __init__(
        self,
        tasks: List[str],
        prompts: LlamaTaskPrompts = LlamaTaskPrompts(),
        task_service_context: Optional[ServiceContext] = None,
    ) -> None:
        self.current_tasks = [Document(x) for x in tasks]
        self.completed_tasks: List[Document] = []
        self.prompts = prompts
        self.task_service_context = task_service_context

    @abstractmethod
    def parse_task_list(self, task_list_str: str) -> List[str]:
        """Parse new tasks generated by the agent."""

    @abstractmethod
    def get_completed_tasks_summary(self) -> str:
        """Generate a summary of completed tasks."""

    @abstractmethod
    def prioritize_tasks(self, objective: str) -> None:
        """Prioritize the current list of incomplete tasks."""

    @abstractmethod
    def generate_new_tasks(
        self, objective: str, prev_task: str, prev_result: str
    ) -> None:
        """Generate new tasks given the previous task and result."""

    @abstractmethod
    def get_next_task(self) -> str:
        """Get the next task to complete."""

    @abstractmethod
    def add_new_tasks(self, tasks: List[str]) -> None:
        """Add new tasks to the task manager."""

    @abstractmethod
    def add_completed_task(self, task: str, result: str) -> None:
        """Add a task as completed."""
