import os
import json
import time

from smolagents import PlanningStep, ActionStep

run_timestamp = str(time.time())

# Storage for logs
_agent_logs: list[dict] = []


logs_path = os.path.join("logs", 'run_logs', f"{run_timestamp}.json")

def extract_tool_calls(tool_calls):
    if not tool_calls:
        return []
    return [{"name": call.name, "args": call.arguments} for call in tool_calls]


def capture_agent_plan_log(step: PlanningStep):
    """
    Accepts agent planning steps & persists them to memory + stdout.
    """

    model_output = getattr(step, 'model_output_message', None)

    output_details = ''

    if model_output:
        output_details = {
            "role": getattr(model_output, 'role', 'assistant'),
            "content": getattr(model_output, 'content', ''),
        }

    log_details = {
        "timestamp": time.time(),
        "plan": str(step.plan),
        "output_details": output_details,
        "tool_calls": extract_tool_calls(getattr(step, 'tool_calls', None)),
    }
    _agent_logs.append(log_details)


def capture_agent_action_log(step: ActionStep):
    """
    Accepts agent action steps & persists them to memory + stdout.
    """

    tool_calls = getattr(step, 'tool_calls', None)
    log_details = {
        "timestamp": time.time(),
        'tool_calls': extract_tool_calls(tool_calls),
        "output": getattr(step, 'output', None),
        'code_action': getattr(step, 'code_action', None),
        'observations': getattr(step, 'observations', None),
        'code_output': getattr(step, 'code_outputs', None),
        'is_final_answer': getattr(step, 'is_final_answer', False),

    }
    _agent_logs.append(log_details)


def flush_logs_to_json():
    """
    Saves all logs for this run to: logs/<run_timestamp>.json
    """
    # Ensure directory exists
    os.makedirs("logs", exist_ok=True)

    print('flushing logs')

    with open(logs_path, "w") as f:
        print(f"Writing {len(_agent_logs)} logs to {logs_path}")
        json.dump(
            {
                "run_id": run_timestamp,
                "logs": _agent_logs,
            },
            f,
            indent=2
        )
    print(f"[AGENT LOG] Saved logs to {logs_path}")