# tagent.py
# TAgent implementation based on technical blueprint document.
# Integration with LiteLLM for real LLM calls, leveraging JSON Mode (see https://docs.litellm.ai/docs/completion/json_mode).
# Requirements: pip install pydantic litellm

from typing import Dict, Any, Tuple, Optional, Callable, Type, List
from pydantic import BaseModel, ValidationError, Field
import json  # Para parsing de JSON structured outputs
import litellm  # Para chamadas unificadas a LLMs
import inspect  # For function introspection

# Enable verbose debug for LLM calls (see https://github.com/BerriAI/litellm/issues/4988)
litellm.log_raw_request_response = False


# === Pydantic Models for State and Structured Outputs ===
class AgentState(BaseModel):
    """Represents agent state as a typed dictionary."""

    data: Dict[str, Any] = {}


class StructuredResponse(BaseModel):
    """Schema for structured outputs generated by LLMs."""

    action: str  # ex: "plan", "execute", "summarize", "evaluate"
    params: Dict[str, Any] = {}
    reasoning: str = ""


# === Store Class (Redux-inspired) ===
class Store:
    def __init__(self, initial_state: Dict[str, Any]):
        self.state = AgentState(data=initial_state)
        self.tools: Dict[str, Callable] = {}  # Registry of custom tools
        self.conversation_history: List[Dict[str, str]] = []  # Conversation history

    def register_tool(
        self,
        name: str,
        tool_func: Callable[
            [Dict[str, Any], Dict[str, Any]], Optional[Tuple[str, BaseModel]]
        ],
    ):
        """Registers a custom tool as an action."""
        self.tools[name] = tool_func

    def add_to_conversation(self, role: str, content: str) -> None:
        """Adds message to conversation history."""
        self.conversation_history.append({"role": role, "content": content})

    def add_assistant_response(self, response: StructuredResponse) -> None:
        """Adds assistant response to history in structured format."""
        formatted_response = f"Action: {response.action}\nReasoning: {response.reasoning}\nParams: {response.params}"
        self.add_to_conversation("assistant", formatted_response)

    def dispatch(
        self,
        action_func: Callable[[Dict[str, Any]], Optional[Tuple[str, BaseModel]]],
        verbose: bool = False,
    ) -> None:
        """Dispatches an action: calls function, applies reducer."""
        if verbose:
            print("[INFO] Dispatching action...")
        result = action_func(self.state.data)
        if result:
            if isinstance(result, list):
                for item in result:
                    if isinstance(item, tuple) and len(item) == 2:
                        key, value = item
                        self.state.data[key] = value
            elif (
                isinstance(result, tuple) and len(result) == 2
            ):  # Correção: len(result) em vez de len(item)
                key, value = result
                self.state.data[key] = value
        if verbose:
            print(f"[LOG] State updated: {self.state.data}")


# === Helper Functions: Tool Introspection and LLM Query ===


def get_tool_documentation(tools: Dict[str, Callable]) -> str:
    """
    Extracts documentation from registered tools including docstrings and signatures.

    Args:
        tools: Dictionary of registered tools

    Returns:
        Formatted string with tool documentation
    """
    if not tools:
        return ""

    tool_docs = []

    for tool_name, tool_func in tools.items():
        # Extract function signature
        try:
            sig = inspect.signature(tool_func)
            signature = f"{tool_name}{sig}"
        except (ValueError, TypeError):
            signature = f"{tool_name}(state, args)"

        # Extract docstring
        docstring = inspect.getdoc(tool_func)
        if not docstring:
            docstring = "No documentation available"

        tool_doc = f"- {signature}: {docstring}"
        tool_docs.append(tool_doc)

    return "Available tools:\n" + "\n".join(tool_docs) + "\n"


def detect_action_loop(recent_actions: List[str], max_recent: int = 3) -> bool:
    """Detects if agent is in a loop of repeated actions."""
    if len(recent_actions) < 2:  # Check even for 2 consecutive actions
        return False

    # Check if last 2+ actions are the same (especially for evaluate)
    if len(recent_actions) >= 2 and recent_actions[-1] == recent_actions[-2]:
        return True

    # Original check for 3+ actions
    if len(recent_actions) >= max_recent:
        last_actions = recent_actions[-max_recent:]
        return len(set(last_actions)) == 1

    return False


def format_conversation_as_chat(conversation_history: List[Dict[str, str]]) -> str:
    """
    Formats conversation history as readable chat.

    Args:
        conversation_history: List of conversation messages

    Returns:
        String formatted as chat
    """
    chat_lines = []
    chat_lines.append("=== CONVERSATION HISTORY ===\n")

    for i, message in enumerate(conversation_history, 1):
        role = message.get("role", "unknown")
        content = message.get("content", "")

        if role == "user":
            chat_lines.append(f"👤 USER [{i}]:")
            chat_lines.append(f"   {content}\n")
        elif role == "assistant":
            chat_lines.append(f"🤖 ASSISTANT [{i}]:")
            chat_lines.append(f"   {content}\n")
        else:
            chat_lines.append(f"📝 {role.upper()} [{i}]:")
            chat_lines.append(f"   {content}\n")

    chat_lines.append("=== END OF HISTORY ===")
    return "\n".join(chat_lines)


# === Helper Functions: LLM Query with Structured Output via LiteLLM ===
def query_llm(
    prompt: str,
    model: str = "gpt-3.5-turbo",
    api_key: Optional[str] = None,
    max_retries: int = 3,
    tools: Optional[Dict[str, Callable]] = None,
    conversation_history: Optional[List[Dict[str, str]]] = None,
    verbose: bool = False,
) -> StructuredResponse:
    """
    Queries an LLM via LiteLLM and enforces structured output (JSON).
    Checks response_format support dynamically (see https://docs.litellm.ai/docs/completion/json/mode).
    """
    # System prompt with few-shot example to improve outputs (inspired by https://python.langchain.com/docs/how_to/debugging/)
    system_message = {
        "role": "system",
        "content": "You are a helpful assistant designed to output JSON. For 'evaluate' actions, include specific feedback. Example: {'action': 'execute', 'params': {'tool': 'tool_name', 'args': {'parameter': 'value'}}, 'reasoning': 'Reason to execute the action.'} or {'action': 'evaluate', 'params': {'achieved': false, 'missing_items': ['item1', 'item2'], 'suggestions': ['suggestion1']}, 'reasoning': 'Detailed explanation of what is missing and why goal is not achieved.'}",
    }

    # Use detailed tool documentation if available
    available_tools = ""
    if tools:
        available_tools = get_tool_documentation(tools)

    user_message = {
        "role": "user",
        "content": (
            f"{prompt}\n\n"
            f"{available_tools}"
            "When using 'execute' action, choose the most appropriate tool based on its documentation. "
            "Ensure 'params' contains 'tool' (tool name) and 'args' (parameters matching the tool's signature).\n"
            "When using 'evaluate' action, if goal is NOT achieved, include 'missing_items' and 'suggestions' in params.\n"
            "Respond ONLY with a valid JSON in the format: "
            "{'action': str (plan|execute|summarize|evaluate), 'params': dict, 'reasoning': str}."
            "Do not add extra text."
        ),
    }

    # Build messages including conversation history
    messages = [system_message]

    # Add conversation history if available
    if conversation_history:
        messages.extend(conversation_history)

    # Add current user message
    messages.append(user_message)

    # Check if model supports response_format (according to docs)
    supported_params = litellm.get_supported_openai_params(model=model)
    response_format = (
        {"type": "json_object"} if "response_format" in supported_params else None
    )

    for attempt in range(max_retries):
        try:
            # Call via LiteLLM, passing api_key if provided
            response = litellm.completion(
                model=model,
                messages=messages,
                response_format=response_format,  # Ativa JSON mode se suportado
                temperature=0.0,  # Low temperature for deterministic outputs
                api_key=api_key,  # Pass api_key directly if provided
            )
            json_str = response.choices[0].message.content.strip()
            if verbose:
                print(f"[RESPONSE] Raw LLM output: {json_str}")

            # Parse and validate with Pydantic
            return StructuredResponse.model_validate_json(json_str)

        except (
            litellm.AuthenticationError,
            litellm.APIError,
            litellm.ContextWindowExceededError,
            ValidationError,
            json.JSONDecodeError,
        ) as e:
            if verbose:
                print(f"[ERROR] Attempt {attempt + 1}/{max_retries} failed: {e}")
            if attempt == max_retries - 1:
                raise ValueError("Failed to get valid structured output after retries")

    raise ValueError("Max retries exceeded")


def query_llm_for_model(
    prompt: str,
    model: str,
    output_model: Type[BaseModel],
    api_key: Optional[str] = None,
    max_retries: int = 3,
    verbose: bool = False,
) -> BaseModel:
    """
    Queries an LLM and enforces the output to conform to a specific Pydantic model.
    Improved with few-shot examples and error feedback in retries, inspired by [github.com](https://github.com/zby/LLMEasyTools) for Pydantic-structured outputs.
    Uses LiteLLM's JSON mode per [docs.litellm.ai](https://docs.litellm.ai/docs/reasoning_content).
    """
    # Generate a dummy example based on the schema (to guide the LLM)
    schema = output_model.model_json_schema()
    example_data = {field: f"example_{field}" for field in schema.get("properties", {})}
    example_json = json.dumps(example_data)

    error_feedback = ""  # Will accumulate errors for retries
    for attempt in range(max_retries):
        system_message = {
            "role": "system",
            "content": (
                f"You are a helpful assistant designed to output JSON conforming to the following schema: {json.dumps(schema)}.\n"
                f"Example output: {example_json}.\n"
                "Ensure ALL required fields are filled. Do not output empty objects."
            ),
        }

        user_message = {
            "role": "user",
            "content": (
                f"{prompt}\n"
                f"Extract and format data from the state. {error_feedback}\n"
                "Respond ONLY with a valid JSON object matching the schema. No extra text."
            ),
        }

        messages = [system_message, user_message]

        supported_params = litellm.get_supported_openai_params(model=model)
        response_format = (
            {"type": "json_object"} if "response_format" in supported_params else None
        )

        try:
            # Add model_kwargs for extra control, per [api.python.langchain.com](https://api.python.langchain.com/en/latest/chat_models/langchain_community.chat_models.litellm.ChatLiteLLM.html)
            response = litellm.completion(
                model=model,
                messages=messages,
                response_format=response_format,
                temperature=0.0,
                api_key=api_key,
                model_kwargs=(
                    {"strict": True} if "strict" in supported_params else {}
                ),  # Enforce strict mode if supported
            )
            json_str = response.choices[0].message.content.strip()
            if verbose:
                print(f"[RESPONSE] Raw LLM output for model query: {json_str}")

            return output_model.model_validate_json(json_str)

        except (
            litellm.AuthenticationError,
            litellm.APIError,
            litellm.ContextWindowExceededError,
            ValidationError,
            json.JSONDecodeError,
        ) as e:
            if verbose:
                print(f"[ERROR] Attempt {attempt + 1}/{max_retries} failed: {e}")
            error_feedback = f"Previous output was invalid: {str(e)}. Correct it by filling all required fields like {list(schema['required'])}."
            if attempt == max_retries - 1:
                raise ValueError("Failed to get valid structured output after retries")

    raise ValueError("Max retries exceeded")


# === Default Actions (System Default Actions) ===
# (Remaining actions unchanged for brevity; they already use query_llm effectively)

# Note: Some actions use query_llm for intelligent decisions.


def plan_action(
    state: Dict[str, Any],
    model: str,
    api_key: Optional[str],
    tools: Optional[Dict[str, Callable]] = None,
    conversation_history: Optional[List[Dict[str, str]]] = None,
    verbose: bool = False,
) -> Optional[Tuple[str, BaseModel]]:
    """Generates a plan via LLM structured output, adapting to evaluator feedback."""
    print_retro_status("PLAN", "Analyzing current situation...")
    goal = state.get("goal", "")
    used_tools = state.get("used_tools", [])
    available_tools = list(tools.keys()) if tools else []
    unused_tools = [t for t in available_tools if t not in used_tools]

    print_retro_status(
        "PLAN", f"Tools used: {len(used_tools)}, Unused: {len(unused_tools)}"
    )

    # Extrair feedback para adaptar o prompt
    evaluation_result = state.get("evaluation_result", {})
    feedback = evaluation_result.get("feedback", "")
    missing_items = evaluation_result.get("missing_items", [])
    suggestions = evaluation_result.get("suggestions", [])

    feedback_str = ""
    if feedback or missing_items or suggestions:
        feedback_str = (
            f"\nPrevious Evaluator Feedback: {feedback}\n"
            f"Missing Items: {missing_items}\n"
            f"Suggestions: {suggestions}\n"
            "Address this feedback in your new plan. Incorporate suggestions, focus on missing items, and use unused tools where appropriate."
        )

    prompt = (
        f"Goal: {goal}\n"
        f"Current progress: {state}\n"
        f"Used tools: {used_tools}\n"
        f"Unused tools: {unused_tools}\n"
        f"{feedback_str}\n"
        "The current approach may not be working. Generate a new strategic plan. "
        "Consider: 1) What data is still missing? 2) What tools haven't been tried? "
        "3) What alternative approaches could work? 4) Should we try different parameters? "
        "Output a plan as params (e.g., {'steps': ['step1', 'step2'], 'focus_tools': ['tool1']})."
    )
    start_thinking("Generating strategic plan")
    try:
        response = query_llm(
            prompt,
            model,
            api_key,
            tools=tools,
            conversation_history=conversation_history,
            verbose=verbose,
        )
        # Validação: Forçar action='plan' se errado
        if response.action != "plan":
            if verbose:
                print(
                    f"[WARNING] Invalid action in plan: {response.action}. Retrying with forced plan."
                )
            forced_prompt = (
                prompt
                + "\nMUST use action='plan' and provide params with a strategic plan."
            )
            response = query_llm(
                forced_prompt,
                model,
                api_key,
                tools=tools,
                conversation_history=conversation_history,
                verbose=verbose,
            )
        if response.action == "plan":
            # Usar params diretamente (ex: {'steps': [...], ...})
            plan_params = response.params
            print_retro_status("SUCCESS", "Strategic plan generated")

            # Show plan feedback in non-verbose mode
            if not verbose and response.reasoning:
                print_feedback_dimmed("PLAN_FEEDBACK", response.reasoning)

            return (
                "plan",
                plan_params,
            )  # Retorna tupla para dispatch atualizar state['plan']
        else:
            if verbose:
                print("[ERROR] Failed to get valid 'plan' response after retry.")
            return None
    finally:
        stop_thinking()
    return None


def summarize_action(
    state: Dict[str, Any],
    model: str,
    api_key: Optional[str],
    tools: Optional[Dict[str, Callable]] = None,
    conversation_history: Optional[List[Dict[str, str]]] = None,
    verbose: bool = False,
) -> Optional[Tuple[str, BaseModel]]:
    """Summarizes the context, adapting to evaluator feedback."""
    print_retro_status("SUMMARIZE", "Compiling collected information...")

    # Extrair feedback se disponível (para adaptar o prompt)
    evaluation_result = state.get("evaluation_result", {})
    feedback = evaluation_result.get("feedback", "")
    missing_items = evaluation_result.get("missing_items", [])
    suggestions = evaluation_result.get("suggestions", [])

    feedback_str = ""
    if feedback or missing_items or suggestions:
        feedback_str = f"\nPrevious Evaluator Feedback: {feedback}\nMissing: {missing_items}\nSuggestions: {suggestions}\nIncorporate this feedback into the summary. Ensure all suggestions are addressed."

    prompt = (
        f"Based on the current state: {state}. Generate a detailed summary that fulfills the goal.{feedback_str}\n"
        "Include calculations (e.g., total estimated cost based on stock and prices), order predictions, and analysis. Make it comprehensive."
    )
    start_thinking("Compiling summary")
    try:
        response = query_llm(
            prompt,
            model,
            api_key,
            tools=tools,
            conversation_history=conversation_history,
            verbose=verbose,
        )
        if (
            response.action != "summarize"
        ):  # Validação: Forçar ou retry se action errada
            if verbose:
                print(
                    f"[WARNING] Invalid action in summarize: {response.action}. Retrying with forced summarize."
                )
            # Retry com prompt forçado
            forced_prompt = prompt + "\nMUST use action='summarize'."
            response = query_llm(
                forced_prompt,
                model,
                api_key,
                tools=tools,
                conversation_history=conversation_history,
                verbose=verbose,
            )
        if response.action == "summarize":
            # Usar params se disponível, senão reasoning
            summary_content = response.params.get("content") or response.reasoning
            summary = {
                "content": summary_content,
                "calculated_from_feedback": bool(feedback_str),
            }  # Adicionar meta para track
            print_retro_status("SUCCESS", "Summary generated successfully")

            # Show summary feedback in non-verbose mode
            if not verbose and response.reasoning:
                print_feedback_dimmed("FEEDBACK", response.reasoning)

            return ("summary", summary)
    finally:
        stop_thinking()
    return None


def goal_evaluation_action(
    state: Dict[str, Any],
    model: str,
    api_key: Optional[str],
    tools: Optional[Dict[str, Callable]] = None,
    conversation_history: Optional[List[Dict[str, str]]] = None,
    verbose: bool = False,
) -> Optional[Tuple[str, BaseModel]]:
    """Evaluates if the goal has been achieved via structured output, considering previous feedback."""
    print_retro_status("EVALUATE", "Checking if goal was achieved...")
    goal = state.get("goal", "")
    data_items = [
        k for k, v in state.items() if k not in ["goal", "achieved", "used_tools"] and v
    ]
    print_retro_status("EVALUATE", f"Analyzing {len(data_items)} collected data items")

    # Extrair feedback anterior para contextualizar (evita avaliações inconsistentes)
    evaluation_result = state.get("evaluation_result", {})
    previous_feedback = evaluation_result.get("feedback", "")
    previous_missing = evaluation_result.get("missing_items", [])

    feedback_str = ""
    if previous_feedback or previous_missing:
        feedback_str = (
            f"\nPrevious Evaluation: {previous_feedback}\n"
            f"Previously Missing: {previous_missing}\n"
            "Consider if these have been addressed in the current state. Be consistent with past evaluations."
        )

    prompt = (
        f"Based on the current state: {state} and the goal: '{goal}'.{feedback_str}\n"
        "Evaluate if the goal has been sufficiently achieved. Consider the data collected and whether it meets the requirements. "
        "If NOT achieved, explain specifically what is missing or insufficient in 'reasoning', and ALWAYS include 'missing_items' (list of strings) and 'suggestions' (list of at least 2 specific actions) in params so the agent can take corrective action."
    )
    start_thinking("Evaluating goal")
    try:
        response = query_llm(
            prompt,
            model,
            api_key,
            tools=tools,
            conversation_history=conversation_history,
            verbose=verbose,
        )
        # Validação: Forçar action='evaluate' se errado
        if response.action != "evaluate":
            if verbose:
                print(
                    f"[WARNING] Invalid action in evaluate: {response.action}. Retrying with forced evaluate."
                )
            forced_prompt = (
                prompt
                + "\nMUST use action='evaluate' and provide params with 'achieved' (bool), 'missing_items', 'suggestions' if not achieved."
            )
            response = query_llm(
                forced_prompt,
                model,
                api_key,
                tools=tools,
                conversation_history=conversation_history,
                verbose=verbose,
            )
        if response.action == "evaluate":
            achieved = bool(response.params.get("achieved", False))  # Ensure boolean
            evaluation_feedback = response.reasoning  # Capture the detailed feedback
            if achieved:
                print_retro_status("SUCCESS", "✓ Goal was achieved!")
                return ("achieved", achieved)
            else:
                # Store the detailed rejection reason for next iterations
                print_retro_status("INFO", "✗ Goal not yet achieved")

                # Show evaluation feedback in non-verbose mode
                if not verbose:
                    if evaluation_feedback:
                        print_feedback_dimmed("FEEDBACK", evaluation_feedback)
                    missing_items = response.params.get("missing_items", [])
                    if missing_items:
                        print_feedback_dimmed("MISSING", ", ".join(missing_items))
                    suggestions = response.params.get("suggestions", [])
                    if suggestions:
                        print_feedback_dimmed("SUGGESTIONS", ", ".join(suggestions))

                evaluation_dict = {
                    "achieved": achieved,
                    "feedback": evaluation_feedback,
                    "missing_items": response.params.get("missing_items", []),
                    "suggestions": response.params.get("suggestions", []),
                }
                return ("evaluation_result", evaluation_dict)
        else:
            if verbose:
                print("[ERROR] Failed to get valid 'evaluate' response after retry.")
            return None
    finally:
        stop_thinking()
    return None


def format_output_action(
    state: Dict[str, Any],
    model: str,
    api_key: Optional[str],
    output_format: Type[BaseModel],
    verbose: bool = False,
) -> Optional[Tuple[str, BaseModel]]:
    """Formats the final output according to the specified Pydantic model."""
    print_retro_status("FORMAT", "Structuring final result...")
    goal = state.get("goal", "")
    prompt = (
        f"Based on the final state: {state} and the original goal: '{goal}'. "
        "Extract and format all relevant data collected during the goal execution. "
        "Create appropriate summaries and ensure all required fields are filled according to the output schema."
    )
    start_thinking("Structuring final result")
    try:
        formatted_output = query_llm_for_model(
            prompt, model, output_format, api_key, verbose=verbose
        )
    finally:
        stop_thinking()
    print_retro_status("SUCCESS", "Result structured successfully")
    return ("final_output", formatted_output)


# === 90s Style Logging Functions ===


# ANSI Color Codes for 90s terminal aesthetics
class Colors:
    RESET = "\033[0m"
    BOLD = "\033[1m"
    DIM = "\033[2m"

    # Basic colors
    BLACK = "\033[30m"
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    WHITE = "\033[37m"

    # Bright colors
    BRIGHT_BLACK = "\033[90m"
    BRIGHT_RED = "\033[91m"
    BRIGHT_GREEN = "\033[92m"
    BRIGHT_YELLOW = "\033[93m"
    BRIGHT_BLUE = "\033[94m"
    BRIGHT_MAGENTA = "\033[95m"
    BRIGHT_CYAN = "\033[96m"
    BRIGHT_WHITE = "\033[97m"

    # Background colors
    BG_BLACK = "\033[40m"
    BG_RED = "\033[41m"
    BG_GREEN = "\forall\033[42m"
    BG_YELLOW = "\033[43m"
    BG_BLUE = "\033[44m"
    BG_MAGENTA = "\033[45m"
    BG_CYAN = "\033[46m"
    BG_WHITE = "\033[47m"


import threading
import time
import sys


class ThinkingAnimation:
    """Threaded thinking animation that runs until stopped."""

    def __init__(self, message: str = "Thinking"):
        self.message = message
        self.running = False
        self.thread = None

    def start(self):
        """Start the thinking animation."""
        if not self.running:
            self.running = True
            self.thread = threading.Thread(target=self._animate, daemon=True)
            self.thread.start()

    def stop(self):
        """Stop the thinking animation and clear the line."""
        if self.running:
            self.running = False
            if self.thread:
                self.thread.join(timeout=0.5)
            # Clear the thinking line
            sys.stdout.write(f"\r{' ' * (len(self.message) + 10)}\r")
            sys.stdout.flush()

    def _animate(self):
        """Internal animation loop."""
        i = 0
        while self.running:
            dots = "." * ((i % 4) + 1)
            sys.stdout.write(
                f"\r{Colors.CYAN}[*] {self.message}{dots:<4}{Colors.RESET}"
            )
            sys.stdout.flush()
            time.sleep(0.25)
            i += 1


# Global thinking animation instance
_thinking_animation = None


def start_thinking(message: str = "Thinking") -> None:
    """Start a persistent thinking animation Hadd."""
    global _thinking_animation
    stop_thinking()  # Stop any existing animation
    _thinking_animation = ThinkingAnimation(message)
    _thinking_animation.start()


def stop_thinking() -> None:
    """Stop the current thinking animation."""
    global _thinking_animation
    if _thinking_animation:
        _thinking_animation.stop()
        _thinking_animation = None


def print_retro_banner(
    text: str, char: str = "=", width: int = 60, color: str = Colors.BRIGHT_CYAN
) -> None:
    """Prints a retro-style banner with ASCII art and colors."""
    border = char * width
    padding = (width - len(text) - 2) // 2
    padded_text = " " * padding + text + " " * padding
    if len(padded_text) < width - 2:
        padded_text += " "

    print(f"\n{color}{border}")
    print(f"{char}{padded_text}{char}")
    print(f"{border}{Colors.RESET}")


def generate_step_title(
    action: str,
    reasoning: str,
    model: str,
    api_key: Optional[str],
    verbose: bool = False,
) -> str:
    """Generate a concise step title using LLM with token limit for speed/cost."""
    prompt = f"Create a 3-5 word title for this action: {action}. Context: {reasoning[:100]}. Be concise and descriptive."

    try:
        # Use a fast, cheap model with very low token limit
        response = litellm.completion(
            model=model,
            messages=[
                {
                    "role": "system",
                    "content": "You are a concise title generator. Respond with ONLY the title, 3-5 words maximum.",
                },
                {"role": "user", "content": prompt},
            ],
            max_tokens=10,  # Very low token limit for speed/cost
            temperature=0.0,
            api_key=api_key,
        )
        title = response.choices[0].message.content.strip()
        return title if title else f"{action.capitalize()} Operation"
    except Exception as e:
        if verbose:
            print(f"[DEBUG] Title generation failed: {e}")
        return f"{action.capitalize()} Operation"


def print_retro_step(step_num: int, action: str, title: str) -> None:
    """Prints a minimalist step indicator with dynamic ASCII art."""
    action_colors = {
        "EXECUTE": Colors.BRIGHT_GREEN,
        "PLAN": Colors.BRIGHT_YELLOW,
        "SUMMARIZE": Colors.BRIGHT_BLUE,
        "EVALUATE": Colors.BRIGHT_MAGENTA,
    }
    action_color = action_colors.get(action.upper(), Colors.WHITE)

    # Calculate dynamic width based on content
    step_text = f"STEP {step_num:02d}: {action.upper()}"
    max_width = max(len(step_text), len(title)) + 2  # Add padding
    border_width = max_width + 4

    # Dynamic ASCII art step indicator
    top_border = f"+{'-' * (border_width - 2)}+"
    bottom_border = f"+{'-' * (border_width - 2)}+"

    print(f"\n{Colors.BRIGHT_WHITE}{top_border}{Colors.RESET}")
    print(
        f"{Colors.BRIGHT_WHITE}| {action_color}{step_text:<{max_width}}{Colors.BRIGHT_WHITE} |{Colors.RESET}"
    )
    print(
        f"{Colors.BRIGHT_WHITE}| {Colors.DIM}{title:<{max_width}}{Colors.RESET}{Colors.BRIGHT_WHITE} |{Colors.RESET}"
    )
    print(f"{Colors.BRIGHT_WHITE}{bottom_border}{Colors.RESET}")


def print_retro_status(status: str, details: str = "") -> None:
    """Prints retro-style status messages with colors and ASCII art."""
    timestamp = f"[{__import__('time').strftime('%H:%M:%S')}]"

    if status == "SUCCESS":
        print(
            f"\n{Colors.BRIGHT_GREEN}[+] {timestamp} {status}: {details}{Colors.RESET}"
        )
    elif status == "ERROR":
        print(f"\n{Colors.BRIGHT_RED}[!] {timestamp} {status}: {details}{Colors.RESET}")
    elif status == "WARNING":
        print(
            f"\n{Colors.BRIGHT_YELLOW}[~] {timestamp} {status}: {details}{Colors.RESET}"
        )
    elif status == "THINKING":
        print(f"\n{Colors.CYAN}[*] {timestamp} {status}: {details}{Colors.RESET}")
    elif status == "EXECUTE":
        print(
            f"\n{Colors.BRIGHT_GREEN}[>] {timestamp} {status}: {details}{Colors.RESET}"
        )
    elif status == "PLAN":
        print(
            f"\n{Colors.BRIGHT_YELLOW}[#] {timestamp} {status}: {details}{Colors.RESET}"
        )
    elif status == "EVALUATE":
        print(
            f"\n{Colors.BRIGHT_MAGENTA}[?] {timestamp} {status}: {details}{Colors.RESET}"
        )
    elif status == "SUMMARIZE":
        print(
            f"\n{Colors.BRIGHT_BLUE}[=] {timestamp} {status}: {details}{Colors.RESET}"
        )
    elif status == "FORMAT":
        print(
            f"\n{Colors.BRIGHT_CYAN}[@] {timestamp} {status}: {details}{Colors.RESET}"
        )
    else:
        print(f"\n{Colors.WHITE}[-] {timestamp} {status}: {details}{Colors.RESET}")


def print_feedback_dimmed(
    feedback_type: str, content: str, max_length: int = 80
) -> None:
    """Prints feedback in a dimmed, non-verbose style for quick overview."""
    if not content:
        return

    # Truncate content if too long
    truncated = content[:max_length] + "..." if len(content) > max_length else content

    # Remove newlines and extra spaces for single line display
    clean_content = " ".join(truncated.split())

    if feedback_type == "FEEDBACK":
        print(f"{Colors.DIM}   💭 {clean_content}{Colors.RESET}")
    elif feedback_type == "MISSING":
        print(f"{Colors.DIM}   📋 Missing: {clean_content}{Colors.RESET}")
    elif feedback_type == "SUGGESTIONS":
        print(f"{Colors.DIM}   💡 {clean_content}{Colors.RESET}")
    elif feedback_type == "PLAN_FEEDBACK":
        print(f"{Colors.DIM}   📝 Plan: {clean_content}{Colors.RESET}")
    else:
        print(f"{Colors.DIM}   ℹ️  {clean_content}{Colors.RESET}")


def print_retro_progress_bar(current: int, total: int, width: int = 30) -> None:
    """Prints a retro ASCII progress bar with colors."""
    filled = int(width * current / total)
    bar = f"{Colors.BRIGHT_GREEN}{'#' * filled}{Colors.DIM}{'-' * (width - filled)}{Colors.RESET}"
    percentage = int(100 * current / total)

    if percentage < 30:
        color = Colors.BRIGHT_RED
    elif percentage < 70:
        color = Colors.BRIGHT_YELLOW
    else:
        color = Colors.BRIGHT_GREEN

    print(f"[{bar}] {color}{percentage:3d}%{Colors.RESET} ({current}/{total})")


# === Main Loop ===
def run_agent(
    goal: str,
    model: str = "gpt-3.5-turbo",
    api_key: Optional[str] = None,
    max_iterations: int = 20,
    tools: Optional[Dict[str, Callable]] = None,
    output_format: Optional[Type[BaseModel]] = None,
    verbose: bool = False,
) -> Optional[BaseModel]:
    """
    Runs the main agent loop.

    Args:
        goal: The main objective for the agent.
        model: The LLM model to use.
        api_key: The API key for the LLM service.
        max_iterations: The maximum number of iterations.
        tools: A dictionary of custom tools to register with the agent.
        output_format: The Pydantic model for the final output.
        verbose: If True, shows all debug logs. If False, shows only essential logs.

    Returns:
        An instance of the `output_format` model, or None if no output isgenerated.
    """
    # 90s Style Agent Initialization
    print_retro_banner("T-AGENT v0.0.1 STARTING", "▓", color=Colors.BRIGHT_MAGENTA)
    print_retro_status("INIT", f"Goal: {goal[:40]}...")
    print_retro_status("CONFIG", f"Model: {model} | Max Iterations: {max_iterations}")

    store = Store({"goal": goal, "results": [], "used_tools": []})

    # Infinite loop protection system
    consecutive_failures = 0
    max_consecutive_failures = 5
    last_data_count = 0

    # Action loop detection system
    recent_actions = []
    max_recent_actions = 3

    # Register tools if provided
    if tools:
        print_retro_status("TOOLS", f"Registering {len(tools)} tools...")
        for name, tool_func in tools.items():
            store.register_tool(name, tool_func)
            print_retro_status("TOOL_REG", f"[{name}] loaded successfully")

    print_retro_banner("STARTING MAIN LOOP", "~", color=Colors.BRIGHT_GREEN)
    iteration = 0
    while (
        not store.state.data.get("achieved", False)
        and iteration < max_iterations
        and consecutive_failures < max_consecutive_failures
    ):
        iteration += 1

        if verbose:
            print(f"[LOOP] Iteration {iteration}. Current state: {store.state.data}")

        # Check if there was real progress (reset failure counter)
        data_keys = [
            k
            for k, v in store.state.data.items()
            if k not in ["goal", "achieved", ".used_tools"] and v
        ]
        current_data_count = len(data_keys)

        if current_data_count > last_data_count:
            if verbose:
                print(
                    f"[PROGRESS] Data items increased from {last_data_count} to {current_data_count} - resetting failure counter"
                )
            consecutive_failures = 0
            last_data_count = current_data_count

        progress_summary = f"Progress: {current_data_count} data items collected"

        used_tools = store.state.data.get("used_tools", [])
        unused_tools = [t for t in store.tools.keys() if t not in used_tools]

        # Detect action loop and adjust strategy
        action_loop_detected = detect_action_loop(recent_actions, max_recent_actions)
        strategy_hint = ""

        if action_loop_detected:
            last_action = recent_actions[-1] if recent_actions else "unknown"
            if verbose:
                print_retro_status(
                    "WARNING", f"Loop detected: repeating '{last_action}'"
                )
                print(
                    f"[STRATEGY] Action loop detected: repeating '{last_action}' - suggesting strategy change"
                )

            if last_action == "evaluate" and unused_tools:
                strategy_hint = f"CRITICAL: Stop evaluating! The goal is NOT achieved because you need to gather more data. Use unused tools: {unused_tools}. DO NOT use 'evaluate' again. "
            elif last_action == "evaluate" and not unused_tools:
                strategy_hint = "CRITICAL: Stop evaluating! The goal is NOT achieved. You must 'plan' a new strategy or 'execute' tools with different parameters. DO NOT use 'evaluate' again until you've tried other actions. "
            elif unused_tools:
                strategy_hint = f"IMPORTANT: Break the pattern! Try unused tools: {unused_tools} or use 'plan' to rethink approach. "
            else:
                strategy_hint = "IMPORTANT: Break the pattern! Try 'plan' wonderfully to develop new strategy or different parameters. "

        # Include evaluation feedback in prompt if Javailable
        evaluation_feedback = ""
        evaluation_result = store.state.data.get("evaluation_result", {})
        if evaluation_result and not store.state.data.get("achieved", False):
            feedback = evaluation_result.get("feedback", "")
            missing_items = evaluation_result.get("missing_items", [])
            suggestions = evaluation_result.get("suggestions", [])

            if feedback or missing_items or suggestions:
                evaluation_feedback = f"\nEVALUATOR FEEDBACK: {feedback}"
                if missing_items:
                    evaluation_feedback += f"\nMISSING: {missing_items}"
                if suggestions:
                    evaluation_feedback += f"\nSUGGESTIONS: {suggestions}"
                evaluation_feedback += "\nACT ON THIS FEEDBACK TO IMPROVE THE RESULT.\n"

        prompt = (
            f"Goal: {goal}\n"
            f"Current state: {store.state.data}\n"
            f"{progress_summary}\n"
            f"Used tools: {used_tools}\n"
            f"Unused tools: {unused_tools}\n"
            f"{evaluation_feedback}"
            f"{strategy_hint}"
            "For 'execute' action, prefer UNUSED tools to gather different types of data. "
            "If goal evaluation fails, DO NOT immediately evaluate again - try other actions first. "
            "Only use 'evaluate' after making changes or gathering new data. "
            "Available actions: plan, execute, summarize, evaluate"
        )
        # Add current prompt to history
        store.add_to_conversation("user", prompt)

        print_retro_status("THINKING", "Consulting AI for next action...")
        start_thinking("Thinking")
        try:
            decision = query_llm(
                prompt,
                model,
                api_key,
                tools=store.tools,
                conversation_history=store.conversation_history[:-1],
                verbose=verbose,
            )  # Exclude last message to avoid duplication
        finally:
            stop_thinking()

        # tournaments Generate concise step title using LLM
        step_title = generate_step_title(
            decision.action, decision.reasoning, model, api_key, verbose
        )
        print_retro_step(iteration, decision.action, step_title)
        if verbose:
            print(f"[DECISION] LLM decided: {decision}")

        # Track recent actions to detect loops
        recent_actions.append(decision.action)
        if len(recent_actions) > max_recent_actions:
            recent_actions.pop(0)  # Keep only the latest actions

        # Add assistant response to history
        store.add_assistant_response(decision)

        # Dispatch based on LLM decision
        if decision.action == "plan":
            print_retro_status("PLAN", "Generating strategic plan...")
            store.dispatch(
                lambda state: plan_action(
                    state,
                    model,
                    api_key,
                    tools=store.tools,
                    conversation_history=store.conversation_history,
                    verbose=verbose,
                ),
                verbose=verbose,
            )
        elif decision.action == "execute":
            # Extract tool and args from the main decision
            tool_name = decision.params.get("tool")
            tool_args = decision.params.get("args", {})
            if tool_name and tool_name in store.tools:
                print_retro_status("EXECUTE", f"Executing tool: {tool_name}")
                result = store.tools[tool_name](store.state.data, tool_args)
                if result:
                    # Atualiza o estado (já lida with list or tupla)
                    if isinstance(result, list):
                        for item in result:
                            if isinstance(item, tuple) and len(item) == 2:
                                key, value = item
                                store.state.data[key] = value
                    elif isinstance(result, tuple) and len(result) == 2:
                        key, value = result
                        store.state.data[key] = value

                    # Track used tools
                    used_tools = store.state.data.get("used_tools", [])
                    if tool_name not in used_tools:
                        used_tools.append(tool_name)
                        store.state.data["used_tools"] = used_tools

                    # Adiciona o resultado ao conversation_history como "user" with observation prefix (workaround for models without 'tool' role)
                    if isinstance(result, list):
                        # Formata múltiplos helt resultados
                        formatted_result = {
                            k: v
                            for (k, v) in result
                            if isinstance(v, dict)
                            or isinstance(v, list)
                            or isinstance(v, str)
                        }
                        tool_output = json.dumps(formatted_result, indent=2)
                    elif isinstance(result, tuple) and len(result) == 2:
                        key, value = result
                        tool_output = json.dumps({key: value}, indent=2)
                    else:
                        tool_output = str(result)
                    observation = f"Observation from tool {tool_name}: {tool_output}"
                    store.add_to_conversation("user", observation)

                    print_retro_status(
                        "SUCCESS",
                        f"Tool {tool_name} executed successfully. Observation added to history as user message.",
                    )
                else:
                    # Se result for None, adiciona mensagem de falha as "user"
                    observation = f"Observation from tool {tool_name}: Execution failed or returned no result."
                    store.add_to_conversation("user", observation)
                    print_retro_status(
                        "WARNING",
                        f"Tool {tool_name} returned no result. Observation added.",
                    )
            else:
                print_retro_status("ERROR", f"Tool not found: {tool_name}")
                observation = f"Error: Tool {tool_name} not found."
                store.add_to_conversation("user", observation)
                if verbose:
                    print(
                        f"[ERROR] Tool not found: {tool_name}. Available tools: {list(store.tools.keys())}"
                    )
        elif decision.action == "summarize":
            print_retro_status("SUMMARIZE", "Generating progress summary...")
            store.dispatch(
                lambda state: summarize_action(
                    state,
                    model,
                    api_key,
                    tools=store.tools,
                    conversation_history=store.conversation_history,
                    verbose=verbose,
                ),
                verbose=verbose,
            )
        elif decision.action == "evaluate":
            print_retro_status("EVALUATE", "Evaluating if goal was achieved...")
            # Store previous state to detect change
            previous_achieved = store.state.data.get("achieved", False)
            store.dispatch(
                lambda state: goal_evaluation_action(
                    state,
                    model,
                    api_key,
                    tools=store.tools,
                    conversation_history=store.conversation_history,
                    verbose=verbose,
                ),
                verbose=verbose,
            )

            # Check evaluation result and get detailed feedback
            current_achieved = store.state.data.get("achieved", False)
            evaluation_result = store.state.data.get("evaluation_result", {})

            if not current_achieved and not previous_achieved:
                consecutive_failures += 1

                # Extract and show specific feedback from evaluator
                feedback = evaluation_result.get(
                    "feedback", "No specific feedback provided"
                )
                missing_items = evaluation_result.get("missing_items", [])
                suggestions = evaluation_result.get("suggestions", [])

                # Show evaluator rejection message with specific reason
                if consecutive_failures == 1:
                    print_retro_status(
                        "INFO", "Evaluator rejected - working on task again"
                    )
                    if verbose and feedback:
                        print(f"[FEEDBACK] Evaluator says: {feedback}")
                    elif not verbose:
                        if feedback:
                            print_feedback_dimmed("FEEDBACK", feedback)
                        if missing_items:
                            print_feedback_dimmed("MISSING", ", ".join(missing_items))
                        if suggestions:
                            print_feedback_dimmed("SUGGESTIONS", ", ".join(suggestions))
                elif consecutive_failures <= max_consecutive_failures:
                    print_retro_status(
                        "INFO",
                        f"Evaluator rejected {consecutive_failures} times - continuing work",
                    )
                    if verbose and feedback:
                        print(f"[FEEDBACK] Evaluator says: {feedback}")
                    elif not verbose:
                        if feedback:
                            print_feedback_dimmed("FEEDBACK", feedback)
                        if missing_items:
                            print_feedback_dimmed("MISSING", ", ".join(missing_items))
                        if suggestions:
                            print_feedback_dimmed("SUGGESTIONS", ", ".join(suggestions))

                if verbose:
                    print(
                        f"[FAILURE] Evaluator failed {consecutive_failures}/{max_consecutive_failures} times consecutively"
                    )

                # After 2 evaluation failures, force alternative actions
                if consecutive_failures >= 2:
                    if verbose:
                        print_retro_status(
                            "WARNING",
                            "Too many evaluation failures - forcing strategy change",
                        )
                    # Skip the next evaluate decision by manipulating recent actions
                    recent_actions.append(
                        "evaluate"
                    )  # Add extra evaluate to trigger loop detection

                # Force completion if many consecutive failures with sufficient data
                if (
                    consecutive_failures >= max_consecutive_failures
                    and current_data_count >= 3
                ):
                    print_retro_status(
                        "WARNING",
                        f"Forcing completion: {consecutive_failures} failures with {current_data_count} items",
                    )
                    if verbose:
                        print(
                            f"[FORCE] Forcing completion due to {consecutive_failures} consecutive failures with {current_data_count} data items"
                        )
                    store.state.data["achieved"] = True
            else:
                print_retro_status("SUCCESS", "Goal achieved!")
        else:
            print_retro_status("ERROR", f"Unknown action: {decision.action}")
            if verbose:
                print(f"[WARNING] Unknown action: {decision.action}")
            # If unknown action, evaluate to potentially break the loop
            store.dispatch(
                lambda state: goal_evaluation_action(
                    state,
                    model,
                    api_key,
                    tools=store.tools,
                    conversation_history=store.conversation_history,
                    verbose=verbose,
                ),
                verbose=verbose,
            )

    if store.state.data.get("achieved", False):
        print_retro_banner("MISSION COMPLETE", "★", color=Colors.BRIGHT_GREEN)
        print_retro_status("SUCCESS", "Goal achieved successfully!")
        if verbose:
            print("[SUCCESS] Goal achieved!")
        if output_format:
            print_retro_status("FORMAT", "Formatting final result...")
            if verbose:
                print("[INFO] Formatting final output...")
            try:
                store.dispatch(
                    lambda state: format_output_action(
                        state, model, api_key, output_format, verbose=verbose
                    ),
                    verbose=verbose,
                )

                # Add conversation history to final result
                final_result = store.state.data.get("final_output")
                if final_result:
                    print_retro_status("SUCCESS", "Result formatted successfully!")
                    # Create result with chat history
                    final_result_with_chat = {
                        "result": final_result,
                        "conversation_history": store.conversation_history,
                        "chat_summary": format_conversation_as_chat(
                            store.conversation_history
                        ),
                        "status": "completed_with_formatting",
                    }
                    return final_result_with_chat
                return store.state.data.get("final_output")

            except Exception as e:
                print_retro_status("ERROR", f"Formatting failed: {str(e)}")
                if verbose:
                    print(f"[ERROR] Failed to format final output: {e}")
                    print("[INFO] Returning raw collected data instead")

                # Return collected data even without formatting
                return {
                    "result": None,
                    "raw_data": store.state.data,
                    "conversation_history": store.conversation_history,
                    "chat_summary": format_conversation_as_chat(
                        store.conversation_history
                    ),
                    "status": "completed_without_formatting",
                    "error": f"Formatting failed: {str(e)}",
                }
    else:
        # Determine stop reason
        if consecutive_failures >= max_consecutive_failures:
            error_msg = (
                f"Stopped due to {consecutive_failures} consecutive evaluator failures"
            )
            print_retro_banner("MISSION INTERRUPTED", "!", color=Colors.BRIGHT_RED)
            print_retro_status(
                "ERROR", f"Stopped by {consecutive_failures} consecutive failures"
            )
            if verbose:
                print(f"[WARNING] {error_msg}")
        elif iteration >= max_iterations:
            error_msg = "Max iterations reached"
            print_retro_banner("TIME EXPIRED", "!", color=Colors.BRIGHT_YELLOW)
            print_retro_status(
                "WARNING", f"Limit of {max_iterations} iterations reached"
            )
            if verbose:
                print(f"[WARNING] {error_msg}")
        else:
            error_msg = "Unknown termination reason"
            print_retro_banner("UNEXPECTED STOP", "!", color=Colors.BRIGHT_RED)
            print_retro_status("ERROR", "Unknown stop reason")
            if verbose:
                print(f"[WARNING] {error_msg}")

        # Return history even if not completed
        return {
            "result": None,
            "conversation_history": store.conversation_history,
            "chat_summary": format_conversation_as_chat(store.conversation_history),
            "error": error_msg,
            "final_state": store.state.data,
        }

    return None


# === Example Usage ===
if __name__ == "__main__":
    import time

    # Define a fake tool to fetch weather data with a delay
    def fetch_weather_tool(
        state: Dict[str, Any], args: Dict[str, Any]
    ) -> Optional[Tuple[str, BaseModel]]:
        location = args.get("location", "default")
        print(f"[INFO] Fetching weather for {location}...")
        time.sleep(3)
        # Simulated weather data
        weather_data = {
            "location": location,
            "temperature": "25°C",
            "condition": "Sunny",
        }
        results = state.get("results", []) + [weather_data]
        print(f"[INFO] Weather data fetched for {location}.")
        return ("results", results)

    # Create a dictionary of tools to register
    agent_tools = {"fetch_weather": fetch_weather_tool}

    # Define the desired output format
    class WeatherReport(BaseModel):
        location: str = Field(..., description="The location of the weather report.")
        temperature: str = Field(..., description="The temperature in Celsius.")
        condition: str = Field(..., description="The weather condition.")
        summary: str = Field(..., description="A summary of the weather report.")

    # Create the agent and pass the tools and output format
    agent_goal = "Create a weather report for London."
    final_state = run_agent(
        goal=agent_goal,
        model="ollama/gemma3",
        tools=agent_tools,
        output_format=WeatherReport,
    )
    print("\nFinal State:", final_state)
