"""Base relentless agent with smart continuation for long tasks."""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from typing import Any

import yaml

from kiss.core import config as config_module
from kiss.core.base import Base
from kiss.core.kiss_agent import KISSAgent
from kiss.core.kiss_error import KISSError
from kiss.core.models.model import Attachment
from kiss.core.printer import Printer
from kiss.docker.docker_manager import DockerManager

TASK_PROMPT = """# Task

{task_description}

# Instructions
- **At step {step_threshold}: you MUST call finish(success=False,
  summary="detailed summary of work done so far")** if the task is
  not complete and you are at risk of running out of steps.
- Work dir: {work_dir}
{previous_progress}
"""

CONTINUATION_PROMPT = """
# Task Progress

{progress_text}

# Continue
- Complete the rest of the task.
- Don't redo completed work.
"""

SUMMARIZER_PROMPT = """
# Summarizer

{trajectory}

# Instructions
- Analyze and return a detailed summary of the work done so far as 'result'.
"""


def finish(success: bool, summary: str) -> str:
    """Finish execution with status and summary.

    Args:
        success: True if successful, False otherwise.
        summary: Detailed summary of work done so far.
    """
    if isinstance(success, str):
        success = success.lower() in ("true", "1", "yes")
    result: str = yaml.dump({"success": bool(success), "summary": summary}, sort_keys=False)
    return result


class RelentlessAgent(Base):
    """Base agent with auto-continuation for long tasks."""

    def _reset(
        self,
        model_name: str | None,
        summarizer_model_name: str | None,
        max_sub_sessions: int | None,
        max_steps: int | None,
        max_budget: float | None,
        work_dir: str | None,
        docker_image: str | None,
        printer: Printer | None = None,
        verbose: bool | None = None,
    ) -> None:
        global_cfg = config_module.DEFAULT_CONFIG
        cfg = global_cfg.assistant.relentless_agent
        default_work_dir = str(Path(global_cfg.agent.artifact_dir).resolve() / "kiss_workdir")

        self.work_dir = str(Path(work_dir or default_work_dir).resolve())
        Path(self.work_dir).mkdir(parents=True, exist_ok=True)

        self.max_sub_sessions = (
            max_sub_sessions if max_sub_sessions is not None else cfg.max_sub_sessions
        )
        self.max_steps = max_steps if max_steps is not None else cfg.max_steps
        self.max_budget = max_budget if max_budget is not None else cfg.max_budget
        self.model_name = model_name if model_name is not None else cfg.model_name
        self.summarizer_model_name = (
            summarizer_model_name
            if summarizer_model_name is not None
            else cfg.summarizer_model_name
        )
        self.budget_used: float = 0.0
        self.total_tokens_used: int = 0
        self.docker_image = docker_image
        self.docker_manager: DockerManager | None = None
        self.set_printer(printer, verbose=verbose)

    def _docker_bash(self, command: str, description: str) -> str:
        if self.docker_manager is None:
            raise KISSError("Docker manager not initialized")
        return self.docker_manager.Bash(command, description)

    def perform_task(
        self,
        tools: list[Callable[..., Any]],
        attachments: list[Attachment] | None = None,
    ) -> str:
        """Execute the task with auto-continuation across multiple sub-sessions.

        Args:
            tools: List of callable tools available to the agent during execution.
            attachments: Optional file attachments (images, PDFs) for the initial prompt.

        Returns:
            YAML string with 'success' and 'summary' keys on successful completion.

        Raises:
            KISSError: If the task fails after exhausting all sub-sessions.
        """
        print(f"Executing task: {self.task_description}")
        all_tools: list[Callable[..., Any]] = [finish, *tools]

        progress_section = ""
        summary = ""
        for trial in range(self.max_sub_sessions):
            executor = KISSAgent(f"{self.name} Trial-{trial}")
            try:
                result = executor.run(
                    model_name=self.model_name,
                    prompt_template=TASK_PROMPT,
                    arguments={
                        "task_description": self.task_description,
                        "previous_progress": progress_section,
                        "step_threshold": str(self.max_steps - 2),
                        "work_dir": self.work_dir,
                    },
                    system_prompt=self.system_instructions,
                    tools=all_tools,
                    max_steps=self.max_steps,
                    max_budget=self.max_budget,
                    printer=self.printer,
                    attachments=attachments if trial == 0 else None,
                )
            except Exception:
                summarizer_agent = KISSAgent(f"{self.name} Summarizer")
                summarizer_result = summarizer_agent.run(
                    model_name=self.summarizer_model_name,
                    prompt_template=SUMMARIZER_PROMPT,
                    is_agentic=False,
                    arguments={
                        "trajectory": executor.get_trajectory(),
                    },
                )
                try:
                    parsed = yaml.safe_load(summarizer_result)
                    summary_text = (
                        parsed.get("result", summarizer_result)
                        if isinstance(parsed, dict)
                        else summarizer_result
                    )
                except Exception:
                    summary_text = summarizer_result
                result = yaml.dump(
                    {"success": False, "summary": summary_text},
                    sort_keys=False,
                )

            self.budget_used += executor.budget_used
            self.total_tokens_used += executor.total_tokens_used

            try:
                payload = yaml.safe_load(result)
            except Exception:
                payload = {}
            if not isinstance(payload, dict):
                payload = {}

            success = payload.get("success", False)
            if isinstance(success, str):
                success = success.lower() in ("true", "1", "yes")
            if success:
                return result

            summary = payload.get("summary", "")
            if summary:
                progress_section = CONTINUATION_PROMPT.format(progress_text=summary)
        raise KISSError(f"Task failed after {self.max_sub_sessions} sub-sessions")

    def run(
        self,
        model_name: str | None = None,
        summarizer_model_name: str | None = None,
        system_instructions: str = "",
        prompt_template: str = "",
        arguments: dict[str, str] | None = None,
        max_steps: int | None = None,
        max_budget: float | None = None,
        work_dir: str | None = None,
        printer: Printer | None = None,
        max_sub_sessions: int | None = None,
        docker_image: str | None = None,
        verbose: bool | None = None,
        tools_factory: Callable[[], list[Callable[..., Any]]] | None = None,
        attachments: list[Attachment] | None = None,
    ) -> str:
        """Run the agent with tools created by tools_factory (called after _reset).

        Args:
            model_name: LLM model to use. Defaults to config value.
            summarizer_model_name: LLM model for summarizing trajectories on failure.
                Defaults to config value.
            system_instructions: System-level instructions passed to the underlying LLM
                via model_config. Defaults to empty string (no system instructions).
            prompt_template: Task prompt template with format placeholders.
            arguments: Dictionary of values to fill prompt_template placeholders.
            max_steps: Maximum steps per sub-session. Defaults to config value.
            max_budget: Maximum budget in USD. Defaults to config value.
            work_dir: Working directory for the agent. Defaults to artifact_dir/kiss_workdir.
            printer: Printer instance for output display.
            max_sub_sessions: Maximum continuation sub-sessions. Defaults to config value.
            docker_image: Docker image name to run tools inside a container.
            verbose: Whether to print output to console. Defaults to config verbose setting.
            tools_factory: Callable that returns the list of tools for the agent.
            attachments: Optional file attachments (images, PDFs) for the initial prompt.

        Returns:
            YAML string with 'success' and 'summary' keys.
        """
        self._reset(
            model_name,
            summarizer_model_name,
            max_sub_sessions,
            max_steps,
            max_budget,
            work_dir,
            docker_image,
            printer,
            verbose,
        )
        self.system_instructions = system_instructions
        self.task_description = prompt_template.format(**(arguments or {}))

        tools = tools_factory() if tools_factory else []

        if self.docker_image:
            with DockerManager(self.docker_image) as docker_mgr:
                self.docker_manager = docker_mgr
                if self.printer:
                    _printer = self.printer

                    def _docker_stream(text: str) -> None:
                        _printer.print(text, type="bash_stream")

                    docker_mgr.stream_callback = _docker_stream
                try:
                    return self.perform_task(tools, attachments=attachments)
                finally:
                    self.docker_manager = None
        return self.perform_task(tools, attachments=attachments)
