import re
import json
from typing import List, Tuple, Optional

from llama_index import Document, ServiceContext
from llama_index.prompts.prompts import QuestionAnswerPrompt, RefinePrompt

from llama_agi.task_manager.base import BaseTaskManager, LlamaTaskPrompts
from llama_agi.utils import initialize_task_list_index
from llama_agi.default_task_prompts import NO_COMPLETED_TASKS_SUMMARY


class LlamaTaskManager(BaseTaskManager):
    """Llama Task Manager

    This task manager uses LlamaIndex to create and prioritize tasks. Using
    the LlamaTaskPrompts, the task manager will create tasks that work
    towards achieving an overall objective.

    New tasks are created based on the prev task+result, completed tasks summary,
    and the overall objective.

    Tasks are then prioritized using the overall objective and current list of tasks.

    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:
        super().__init__(
            tasks=tasks, prompts=prompts, task_service_context=task_service_context
        )

        self.current_tasks_index = initialize_task_list_index(
            self.current_tasks, service_context=self.task_service_context
        )
        self.completed_tasks_index = initialize_task_list_index(
            self.completed_tasks, service_context=self.task_service_context
        )

        self.task_create_qa_template = self.prompts.task_create_qa_template
        self.task_create_refine_template = self.prompts.task_create_refine_template

        self.task_prioritize_qa_template = self.prompts.task_prioritize_qa_template
        self.task_prioritize_refine_template = (
            self.prompts.task_prioritize_refine_template
        )

    def _get_task_create_templates(
        self, prev_task: str, prev_result: str
    ) -> Tuple[QuestionAnswerPrompt, RefinePrompt]:
        """Fetch the task create prompts as llama_index objects."""
        text_qa_template = self.task_create_qa_template.format(
            prev_result=prev_result,
            prev_task=prev_task,
            query_str="{query_str}",
            context_str="{context_str}",
        )
        llama_text_qa_template = QuestionAnswerPrompt(text_qa_template)

        refine_template = self.task_create_refine_template.format(
            prev_result=prev_result,
            prev_task=prev_task,
            query_str="{query_str}",
            context_msg="{context_msg}",
            existing_answer="{existing_answer}",
        )
        llama_refine_template = RefinePrompt(refine_template)

        return (llama_text_qa_template, llama_refine_template)

    def _get_task_prioritize_templates(
        self,
    ) -> Tuple[QuestionAnswerPrompt, RefinePrompt]:
        """Fetch the task prioritize prompts as llama_index objects."""
        return (
            QuestionAnswerPrompt(self.task_prioritize_qa_template),
            RefinePrompt(self.task_prioritize_refine_template),
        )

    def parse_task_list(self, task_list_str: str) -> List[str]:
        """Parse new tasks generated by the agent."""
        new_tasks: List[str] = []
        try:
            new_tasks = json.loads(task_list_str)
            new_tasks = [x.strip() for x in new_tasks if len(x.strip()) > 10]
        except Exception:
            new_tasks = str(task_list_str).split("\n")
            new_tasks = [
                re.sub(r"^[0-9]+\.", "", x).strip()
                for x in str(new_tasks)
                if len(x.strip()) > 10 and x[0].isnumeric()
            ]
        return new_tasks

    def get_completed_tasks_summary(self) -> str:
        """Generate a summary of completed tasks."""
        if len(self.completed_tasks) == 0:
            return NO_COMPLETED_TASKS_SUMMARY
        summary = self.completed_tasks_index.as_query_engine(
            response_mode="tree_summarize"
        ).query(
            "Summarize the current completed tasks",
        )
        return str(summary)

    def prioritize_tasks(self, objective: str) -> None:
        """Prioritize the current list of incomplete tasks."""
        (text_qa_template, refine_template) = self._get_task_prioritize_templates()
        prioritized_tasks = self.current_tasks_index.as_query_engine(
            text_qa_template=text_qa_template, refine_template=refine_template
        ).query(objective)

        new_tasks = []
        for task in str(prioritized_tasks).split("\n"):
            task = re.sub(r"^[0-9]+\.", "", task).strip()
            if len(task) > 10:
                new_tasks.append(task)
        self.current_tasks = [Document(x) for x in new_tasks]
        self.current_tasks_index = initialize_task_list_index(
            self.current_tasks, service_context=self.task_service_context
        )

    def generate_new_tasks(
        self, objective: str, prev_task: str, prev_result: str
    ) -> None:
        """Generate new tasks given the previous task and result."""
        (text_qa_template, refine_template) = self._get_task_create_templates(
            prev_task, prev_result
        )
        task_list_response = self.completed_tasks_index.as_query_engine(
            text_qa_template=text_qa_template, refine_template=refine_template
        ).query(objective)
        new_tasks = self.parse_task_list(str(task_list_response))
        self.add_new_tasks(new_tasks)

    def get_next_task(self) -> str:
        """Get the next task to complete."""
        next_task = self.current_tasks.pop().get_text()
        self.current_tasks_index = initialize_task_list_index(
            self.current_tasks, service_context=self.task_service_context
        )
        return next_task

    def add_new_tasks(self, tasks: List[str]) -> None:
        """Add new tasks to the task manager."""
        for task in tasks:
            if task not in self.current_tasks and task not in self.completed_tasks:
                self.current_tasks.append(Document(task))
        self.current_tasks_index = initialize_task_list_index(
            self.current_tasks, service_context=self.task_service_context
        )

    def add_completed_task(self, task: str, result: str) -> None:
        """Add a task as completed."""
        document = Document(f"Task: {task}\nResult: {result}\n")
        self.completed_tasks.append(document)
        self.completed_tasks_index = initialize_task_list_index(
            self.completed_tasks, service_context=self.task_service_context
        )
