# -*- coding: utf-8 -*-
# pylint: disable=E0213
"""
Data structures about the roadmap for complicated tasks
"""
from datetime import datetime
from typing import List, Literal, Tuple, Optional, Any, Dict
from pydantic import BaseModel, Field, field_validator


def get_current_time_message() -> str:
    """
    Returns the current time as a formatted string.

    Returns:
        str: The current time formatted as 'YYYY-MM-DD HH:MM:SS'.
    """
    return f"Current time is {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"


WORKER_PROGRESS_SUMMARY = (
    "## Instruction\n"
    "Review the execution trace above and generate a comprehensive summary "
    "report in Markdown format that addresses the original task/query. "
    "Your report must include:\n\n"
    "1. **Task Overview**\n"
    "   - Include the original query/task verbatim;\n"
    "   - Briefly state the main objective.\n"
    "2. **Comprehensive Analysis**"
    "   - Provide a detailed, structured answer to the original query/task;\n"
    "   - Include all relevant information requested in the original task;\n"
    "   - Support your findings with specific references from your execution "
    "trace;\n"
    "   - Organize content into logical sections with appropriate headings;\n"
    "   - Include data visualizations, tables, or formatted lists when "
    "applicable.\n\n"
    "3. **Completion Checklist**\n"
    "   - Reproduce the original 'Expected Output' checklist of required "
    "tasks/information; **NEVER** makeup additional expected output items "
    "in the checklist\n"
    "   - Mark each item as [x] Completed or [ ] Incomplete;\n"
    "   - For each completed item, reference where in your report this "
    "information appears;\n"
    "   - For incomplete items, explain briefly why they remain unaddressed;\n"
    "4. **Conclusion**\n"
    "   - If the task is fully complete, provide a brief conclusion "
    "summarizing key findings;\n"
    "   - If the task remains incomplete, outline a specific plan to "
    "address remaining items, including:\n"
    "     - Which tools would be used;\n"
    "     - What information is still needed;\n"
    "     - Sequence of planned actions.\n\n"
    "Format your report professionally with consistent heading levels, "
    "proper spacing, and appropriate emphasis for key information."
)


WORKER_NEXT_STEP_INSTRUCTION = """
If the subtask remains incomplete, outline a specific plan to address remaining
items, including:
     - Which tools would be used
     - What information is still needed
     - Sequence of planned actions
Leave it as an empty string is the subtask has been done successfully.
"""

WORKER_FILE_COLLECTION_INSTRUCTION = (
    "Collect all files generated in the execution process, "
    "such as the files generated by `write_file` and `edit_file`."
    "This field MUST be in dictionary, where"
    "the keys are the paths of generated files "
    "(e.g. '/FULL/PATH/OF/FILE_1.md') and the values are short "
    "descriptions about the generated files."
)


class WorkerResponse(BaseModel):
    """
    Represents the response structure from a worker agent after task execution.

    This class defines the expected format for worker responses, including
    progress summaries, next steps, tool usage information, and task
    completion status.

    Attributes:
        subtask_progress_summary (str):
            Comprehensive summary report of task execution.
        next_step (str):
            Description of planned next actions if task is incomplete.
        generated_files (dict):
            Dictionary mapping file paths to descriptions of generated files.
        task_done (bool):
            Flag indicating whether the task has been completed.
    """

    subtask_progress_summary: str = Field(
        ...,
        description=WORKER_PROGRESS_SUMMARY,
    )
    next_step: str = Field(
        ...,
        description=WORKER_NEXT_STEP_INSTRUCTION,
    )
    generated_files: dict = Field(
        ...,
        description=WORKER_FILE_COLLECTION_INSTRUCTION,
    )
    task_done: bool = Field(
        ...,
        description="Whether task is done or it require addition effort",
    )


class Update(BaseModel):
    """Represents an update record from a worker during task execution.

    This class tracks progress updates from workers as they work on subtasks,
    including status changes, progress summaries, and execution details.

    Attributes:
        reason_for_status (str): Explanation for the current status.
        task_done (bool): Whether the task has been completed.
        subtask_progress_summary (str): Summary of progress made.
        next_step (str): Description of planned next actions.
        worker (str): Identifier of the worker providing the update.
        attempt_idx (int): Index of the current attempt.
    """

    reason_for_status: str
    task_done: bool
    subtask_progress_summary: str
    next_step: str
    worker: str
    attempt_idx: int

    @field_validator(
        "subtask_progress_summary",
        "reason_for_status",
        "next_step",
        "worker",
        mode="before",
    )
    def _stringify(cls, v: Any) -> str:
        """ensure the attributes are string"""
        if v is None:
            return ""
        return str(v)


class WorkerInfo(BaseModel):
    """Contains information about a worker agent assigned to a subtask.

    This class stores metadata about worker agents, including their
    capabilities, creation type, and configuration details.

    Attributes:
        worker_name (str):
            Name identifier of the worker.
        status (str):
            Current status of the worker.
        create_type (Literal["built-in", "dynamic-built"]):
            How the worker was created.
        description (str):
            Description of the worker's purpose and capabilities.
        tool_lists (List[str]):
            List of tools available to this worker.
        sys_prompt (str):
            System prompt used to configure the worker.
    """

    worker_name: str = ""
    status: str = ""
    create_type: Literal["built-in", "dynamic-built"] = "dynamic-built"
    description: str = ""
    # for dynamically create worker agents
    tool_lists: List[str] = Field(default_factory=list)
    sys_prompt: str = ""

    @field_validator(
        "worker_name",
        "status",
        mode="before",
    )
    def _stringify(cls, v: Any) -> str:
        if v is None:
            return ""
        return str(v)


class SubTaskSpecification(BaseModel):
    """
    Details of a subtask within a larger task decomposition.
    """

    subtask_description: str = Field(description="Description of the subtask.")
    input_intro: str = Field(
        ...,
        description="Introduction or context for the subtask input.",
    )
    exact_input: str = Field(
        ...,
        description="The exact input data or parameters for the subtask.",
    )
    expected_output: str = Field(
        ...,
        description="The expected output data or parameters for the subtask.",
    )
    desired_auxiliary_tools: str = Field(
        ...,
        description="Tools that would be helpful for this subtask.",
    )

    @field_validator(
        "subtask_description",
        "input_intro",
        "exact_input",
        "expected_output",
        "desired_auxiliary_tools",
        mode="before",
    )
    def _stringify(cls, v: Any) -> str:
        if v is None:
            return ""
        return str(v)


class SubTaskStatus(BaseModel):
    """
    Represents the status and details of a subtask within a
    larger task decomposition.

    This class tracks individual subtasks, their execution status,
    assigned workers, and progress updates throughout the execution lifecycle.

    Attributes:
        status (Literal["Planned", "In-process", "Done"]):
            Current execution status.
        updates (List[Update]):
            List of progress updates from workers.
        attempt (int):
            Number of execution attempts for this subtask.
        workers (List[WorkerInfo]):
            List of workers assigned to this subtask.
    """

    subtask_specification: SubTaskSpecification = Field(
        default_factory=SubTaskSpecification,
    )
    status: Literal["Planned", "In-process", "Done"] = "Planned"
    updates: List[Update] = Field(
        default_factory=list,
        description=(
            "List of updates from workers. "
            "MUST be empty list when initialized."
        ),
    )
    attempt: int = 0
    workers: List[WorkerInfo] = Field(
        default_factory=list,
        description=(
            "List of workers that have been assigned to his subtask."
            "MUST be EMPTY when initialize the subtask."
        ),
    )


class RoadMap(BaseModel):
    """Represents a roadmap for task decomposition and execution tracking.

    This class manages the overall task breakdown, containing the original task
    description and a list of decomposed subtasks with their execution status.

    Attributes:
        original_task (str):
            The original task description before decomposition.
        decomposed_tasks (List[SubTaskStatus]):
            List of subtasks created from the original task.
    """

    original_task: str = ""
    decomposed_tasks: List[SubTaskStatus] = Field(default_factory=list)

    def next_unfinished_subtask(
        self,
    ) -> Tuple[Optional[int], Optional[SubTaskStatus]]:
        """Find the next subtask that is not yet completed.

        Iterates through the decomposed tasks to find the first subtask
        with status "Planned" or "In-process".

        Returns:
            Tuple[Optional[int], Optional[SubTaskStatus]]: A tuple containing:
                - The index of the next unfinished subtask
                    (None if all tasks are done)
                - The SubTaskStatus object of the next unfinished subtask
                    (None if all tasks are done)
        """
        for i, subtask in enumerate(self.decomposed_tasks):
            if subtask.status in ["Planned", "In-process"]:
                return i, subtask
        return None, None


class PlannerNoteBook(BaseModel):
    """
    Represents a planner notebook.

    Attributes:
        time (str): The current time message.
        user_input (List[str]): List of user inputs.
        detail_analysis_for_plan (str): Detailed analysis for the plan.
        roadmap (RoadMap): The roadmap associated with the planner.
        files (Dict[str, str]): Dictionary of files related to the planner.
        full_tool_list (dict[str, dict]): Full schema of tools.
    """

    time: str = Field(default_factory=get_current_time_message)
    user_input: List[str] = Field(default_factory=list)
    detail_analysis_for_plan: str = (
        "Unknown. Please call `build_roadmap_and_decompose_task` to analyze."
    )
    roadmap: RoadMap = Field(default_factory=RoadMap)
    files: Dict[str, str] = Field(default_factory=dict)
    full_tool_list: list[dict] = Field(default_factory=list)
