# Author: Koushik Sen (ksen@berkeley.edu)
# Contributors:
# Koushik Sen (ksen@berkeley.edu)
# add your name here

"""Base agent class with common functionality for all KISS agents."""

import json
import logging
import sys
import threading
import time
from pathlib import Path
from typing import Any, ClassVar

import yaml
from yaml.nodes import ScalarNode

from kiss.core import config as config_module
from kiss.core.models.model_info import get_max_context_length
from kiss.core.print_to_console import ConsolePrinter
from kiss.core.printer import Printer
from kiss.core.utils import config_to_dict

logger = logging.getLogger(__name__)

def _str_presenter(dumper: yaml.Dumper, data: str) -> ScalarNode:
    """Use literal block style for multiline strings in YAML output."""
    return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")  # type: ignore[reportUnknownMemberType]


yaml.add_representer(str, _str_presenter)

_KISS_DIR = Path().home() / ".kiss"

_artifact_dir = Path(config_module.DEFAULT_CONFIG.agent.artifact_dir)

SYSTEM_PROMPT = f"""
# FOCUS ON THE GIVEN TASK.  IT'S COMPLETION IS YOUR SOLE GOAL.  BE RELENTLESS.

# Rules
- Write() for new files. Edit() for small changes.
- Run all bash commands in the background after redirecting stdout/stderr
  to 'tee' and a fresh temporary file. Poll the tail of the temporary file
  every 5 seconds to check progress of the bash command.
- Use go_to_url() for browser tool and internet search or testing an agent/app.
- Look at `{_KISS_DIR}/task_history.jsonl` for task history and context.  Recent tasks
  are appended to the file.  The file could be very large.
  Pay more attention to the recent tasks over old tasks. Do not try to finish a
  task from the task history. DO NOT WRITE/EDIT the task history.
- If you don't know the context of a vague task, look at the latest tasks in the
  task history from latest to oldest to get the context.
- Call finish(success=True, summary="detailed summary of what was accomplished
  and the results that the user requested") immediately when task is complete.
- Whenever the user asks the agent to show something, create a webpage and show
  the webpage in the user's default browser.
- Use 'uv run myprogram.py' for running Python programs.
- READ large files in chunks.
- Create temporary files in {_artifact_dir.parent}/tmp
- YOU **MUST FOLLOW THE INSTRUCTIONS DIRECTLY**

## Code Style Guidelines
- Write simple, clean, and readable code with minimal indirection
- Avoid unnecessary object attributes, local variables, and config variables
- Avoid tight coupling between files and modules.
- Avoid object/struct attribute redirections
- No redundant abstractions or duplicate code
- Each function should do one thing well
- Use clear, descriptive names
- Public methods MUST have full documentation
- Understand the root cause of an issue or bug, and patch the root cause instead of
  of an ad hoc superficial fix.
- Before you write code, wait and think if the code is simple, elegant, general, and minimal.
- Once you finish the task, DO NOT write documentations unless the task specifically requires it.
- You MUST check and test the code you have written except for formatting/typing changes

## Testing Instructions
- Run lint and typecheckers and fix any lint and typecheck errors.
  Use 'uv run check --full' if available.
- Carefully read the code, find and fix redundancies, duplications,
  inconsistencies, errors, and AI slop in the code
- You MUST achieve 100% branch coverage
- Tests MUST NOT use mocks, patches, fakes, or any form of test doubles
- Integration tests are HIGHLY encouraged
- You MUST not add tests that are redundant or duplicate of existing
  tests or does not add new branch coverage over existing tests
- Generate meaningful stress tests for the code if you are
  optimizing the code for performance
- Each test should be independent and verify actual behavior

## Use web tools when you need to:
- Look up API documentation or library usage from the internet
- Find examples of similar implementations
- Understand existing code in the project
- Read papers from the internet to understand concepts and algorithms
- For deep research, you must visit and read at least 50 websites

### Self-Improvement Loop
- Just before finishing an agent task, update `{_artifact_dir.parent}/LESSONS.md`
  with instructions and rules and intelligence for yourself ONLY IF you have learned any
  major lessons or intelligence about the project or in general during the task execution.
  Lessons that save running time and number of tokens used by the agent would be invaluable.
  Also compact the lessons you have learned into concise instructions and rules for you to follow.
- The lessons MUST NOT be specific to a task, but about agent behavior.
- Review lessons when the agent starts

## After you have implemented the task, aggresively and carefully simplify and clean up the code
 - Remove unnecessary conditional checks
 - Make sure that the code is still working correctly
 - Simplify and clean up the test code
 - Remove all temporary files you created
"""


class Base:
    """Base class for all KISS agents with common state management and persistence."""

    agent_counter: ClassVar[int] = 1
    global_budget_used: ClassVar[float] = 0.0
    _class_lock: ClassVar[threading.Lock] = threading.Lock()

    model_name: str
    messages: list[dict[str, Any]]
    function_map: Any
    run_start_timestamp: int
    budget_used: float
    total_tokens_used: int
    step_count: int
    printer: Printer | None

    def __init__(self, name: str) -> None:
        """Initialize a Base agent instance.

        Args:
            name: The name identifier for the agent.
        """
        self.name = name
        with Base._class_lock:
            self.id = Base.agent_counter
            Base.agent_counter += 1
        self.base_dir = ""
        self.printer: Printer | None = None

    def set_printer(
        self,
        printer: Printer | None = None,
        verbose: bool | None = None,
    ) -> None:
        """Configure the output printer for this agent.

        If an explicit *printer* is provided, it is always used regardless
        of the verbose setting.  Otherwise a ``ConsolePrinter`` is created
        when verbose output is enabled (either explicitly or via config).

        Args:
            printer: An existing Printer instance to use directly. If provided,
                verbose is ignored.
            verbose: Whether to print to the console. If None,
                uses the verbose config value.
        """
        if printer:
            self.printer = printer
        elif verbose is not False and config_module.DEFAULT_CONFIG.agent.verbose:
            self.printer = ConsolePrinter()
        else:
            self.printer = None

    def _build_state_dict(self) -> dict[str, Any]:
        """Build state dictionary for saving.

        Returns:
            dict[str, Any]: A dictionary containing all agent state for persistence.
        """
        try:
            max_tokens = get_max_context_length(self.model_name)
        except Exception:
            logger.debug("Exception caught", exc_info=True)
            max_tokens = None

        return {
            "name": self.name,
            "id": self.id,
            "messages": self.messages,
            "function_map": list(self.function_map),
            "run_start_timestamp": self.run_start_timestamp,
            "run_end_timestamp": int(time.time()),
            "config": config_to_dict(),
            "arguments": getattr(self, "arguments", {}),
            "prompt_template": getattr(self, "prompt_template", ""),
            "is_agentic": getattr(self, "is_agentic", True),
            "model": self.model_name,
            "budget_used": self.budget_used,
            "total_budget": getattr(
                self, "max_budget", config_module.DEFAULT_CONFIG.agent.max_agent_budget
            ),
            "global_budget_used": Base.global_budget_used,
            "global_max_budget": config_module.DEFAULT_CONFIG.agent.global_max_budget,
            "tokens_used": self.total_tokens_used,
            "max_tokens": max_tokens,
            "step_count": self.step_count,
            "max_steps": getattr(self, "max_steps", config_module.DEFAULT_CONFIG.agent.max_steps),
            "command": " ".join(sys.argv),
        }

    def _save(self) -> None:
        """Save the agent's state to a YAML file in the artifacts directory.

        The file is saved to {artifact_dir}/trajectories/trajectory_{name}_{id}_{timestamp}.yaml
        """
        folder_path = Path(config_module.DEFAULT_CONFIG.agent.artifact_dir) / "trajectories"
        folder_path.mkdir(parents=True, exist_ok=True)
        name_safe = self.name.replace(" ", "_").replace("/", "_")
        filename = folder_path / f"trajectory_{name_safe}_{self.id}_{self.run_start_timestamp}.yaml"
        with filename.open("w", encoding="utf-8") as f:
            yaml.dump(self._build_state_dict(), f, indent=2)

    def get_trajectory(self) -> str:
        """Return the trajectory as JSON for visualization.

        Returns:
            str: A JSON-formatted string of all messages in the agent's history.
        """
        return json.dumps(self.messages, indent=2)

    def _add_message(self, role: str, content: Any, timestamp: int | None = None) -> None:
        """Add a message to the history.

        Args:
            role: The role of the message sender (e.g., 'user', 'model').
            content: The content of the message.
            timestamp: Optional Unix timestamp. If None, uses current time.
        """
        self.messages.append(
            {
                "unique_id": len(self.messages),
                "role": role,
                "content": content,
                "timestamp": timestamp if timestamp is not None else int(time.time()),
            }
        )
