import json

from smolagents import MultiStepAgent
from smolagents.memory import ActionStep, FinalAnswerStep, PlanningStep


def _format_arguments(arguments) -> str:
    if arguments is None:
        return "{}"
    if isinstance(arguments, str):
        return arguments
    try:
        return json.dumps(arguments, ensure_ascii=False, indent=2)
    except TypeError:
        return str(arguments)


def _truncate(text: str, limit: int = 600) -> str:
    if not text:
        return "(пусто)"
    text = str(text).strip()
    if len(text) <= limit:
        return text
    return text[:limit] + f"\n… (ещё {len(text) - limit} символов)"


class StreamingAgentRunner:
    def __init__(self, agent: MultiStepAgent):
        self.agent = agent

    def _format_action_step(self, step: ActionStep) -> list[str]:
        lines = [f"### Шаг {step.step_number}"]

        if step.model_output:
            output = step.model_output
            if isinstance(output, list):
                output = json.dumps(output, ensure_ascii=False)
            lines.append("**Модель:**")
            lines.append(f"```\n{_truncate(output, 400)}\n```")

        if step.tool_calls:
            for i, tc in enumerate(step.tool_calls, start=1):
                prefix = f"[{step.step_number}.{i}]" if len(step.tool_calls) > 1 else f"[{step.step_number}]"
                lines.append(f"**{prefix} → `{tc.name}`**")
                lines.append(f"```json\n{_format_arguments(tc.arguments)}\n```")

        if step.observations:
            lines.append("**← результат:**")
            lines.append(f"```\n{_truncate(step.observations)}\n```")

        if step.error:
            lines.append(f"**⚠ ошибка:** {step.error}")

        return lines

    def run(self, query: str):
        """Генератор прогресса агента и финального ответа."""
        parts = ["Агент думает...\n"]
        yield parts[0], 0

        try:
            for step in self.agent.run(query, stream=True, reset=True):
                if isinstance(step, PlanningStep):
                    parts.append(f"**План:** {step.plan.strip()}")
                elif isinstance(step, ActionStep):
                    parts.extend(self._format_action_step(step))
                elif isinstance(step, FinalAnswerStep):
                    parts.append(f"## Ответ\n\n{step.output}")

                progress = min(len(parts) / (self.agent.max_steps + 1), 0.95)
                yield "\n\n".join(parts), progress

            yield "\n\n".join(parts), 1.0

        except Exception as e:
            yield f"Ошибка: {e}", 1.0