# ============================================================================
# agent_py.py — Simple Python Agent backend (raw Mistral SDK tool-calling loop)
# ============================================================================
#
# CONTRACT: BACKEND_NAME, get_client, run, build_code_snippets
#
# PATTERN
# -------
# Classic tool-calling loop. The LLM sees the user's message plus a list
# of tool schemas. On each iteration it either:
#   - emits tool calls (we run them and append results to the history), or
#   - emits plain text (loop exits with that as the final reply).
#
# Bounded by MAX_AGENT_STEPS. No framework. Pure Python against the raw
# Mistral SDK.
# ============================================================================

import os
import json

# Defensive import — see agent_workflow.py for full explanation.
_Mistral = None
try:
    from mistralai import Mistral as _Mistral  # v1.x
except ImportError:
    try:
        from mistralai.client import Mistral as _Mistral  # v2.x
    except ImportError:
        try:
            from mistralai.client import MistralClient as _OldClient  # v0.x
            from mistralai.models.chat_completion import ChatMessage as _OldMsg

            class _ChatShim:
                def __init__(self, client):
                    self._client = client
                def complete(self, model, messages, temperature=None,
                             max_tokens=None, tools=None):
                    msgs = [_OldMsg(role=m["role"], content=m.get("content", ""))
                            for m in messages]
                    return self._client.chat(
                        model=model, messages=msgs,
                        temperature=temperature, max_tokens=max_tokens,
                        tools=tools,
                    )

            class _MistralV0Wrapper:
                def __init__(self, api_key):
                    self._client = _OldClient(api_key=api_key)
                    self.chat = _ChatShim(self._client)

            _Mistral = _MistralV0Wrapper
        except ImportError as _e:
            raise ImportError(
                "mistralai package is missing or an unknown version. "
                f"Last error: {_e}"
            )

Mistral = _Mistral

from parameters import TEMPERATURE, MAX_TOKENS, MAX_AGENT_STEPS
from prompts import AGENT_SYSTEM
from tools import TOOL_FUNCTIONS, TOOL_SCHEMAS
import providers


BACKEND_NAME = "Simple Python Agent"


def get_client(api_key, provider="Mistral"):
    return providers.get_llm_client(provider, api_key)


def _llm(client, messages, tools=None, provider="Mistral"):
    model = providers.get_llm_model(provider)
    return client.chat.complete(
        model=model,
        temperature=TEMPERATURE,
        max_tokens=MAX_TOKENS,
        messages=messages,
        tools=tools,
    ).choices[0].message


def run(client, user_message, provider="Mistral"):
    """Tool-calling loop. LLM decides when to stop."""
    messages = [
        {"role": "system", "content": AGENT_SYSTEM},
        {"role": "user", "content": user_message},
    ]
    steps = []
    tool_calls_made = []

    for step_num in range(1, MAX_AGENT_STEPS + 1):
        msg = _llm(client, messages, tools=TOOL_SCHEMAS, provider=provider)
        messages.append(providers.serialize_assistant_message(msg, provider))
        tool_calls = msg.tool_calls or []

        if not tool_calls:
            # No more tool calls — model has a final answer
            steps.append({
                "step": step_num, "type": "final", "tool": "-",
                "args": "-", "result": msg.content or "",
            })
            return {
                "reply": msg.content or "",
                "steps": steps,
                "extracted": {"tool_calls_made": tool_calls_made},
            }

        # Execute each tool call and append results
        for tc in tool_calls:
            name = tc.function.name
            args_raw = tc.function.arguments
            args = json.loads(args_raw) if isinstance(args_raw, str) else args_raw
            result = TOOL_FUNCTIONS[name](**args)
            steps.append({
                "step": step_num, "type": "tool_call", "tool": name,
                "args": json.dumps(args), "result": str(result),
            })
            tool_calls_made.append({"tool": name, "args": args, "result": result})
            messages.append(providers.serialize_tool_result(tc, name, result, provider))

    steps.append({
        "step": MAX_AGENT_STEPS, "type": "limit", "tool": "-",
        "args": "-", "result": "max steps reached",
    })
    return {
        "reply": "(max agent steps reached)",
        "steps": steps,
        "extracted": {"tool_calls_made": tool_calls_made},
    }


def build_code_snippets(user_message, steps):
    lines = [
        "# Backend: Simple Python Agent",
        "# Raw Mistral SDK tool-calling loop. No framework.",
        f"# User message: {user_message}",
        "",
        "messages = [",
        "    {'role': 'system', 'content': AGENT_SYSTEM},",
        f"    {{'role': 'user', 'content': {user_message!r}}},",
        "]",
        "",
        "for step in range(1, MAX_AGENT_STEPS + 1):",
        "    msg = client.chat.complete(",
        "        model=MODEL, messages=messages, tools=TOOL_SCHEMAS",
        "    ).choices[0].message",
        "    messages.append(msg.model_dump(exclude_none=True))",
        "",
        "    if not msg.tool_calls:",
        "        break  # plain-text reply means we are done",
        "",
        "    for tc in msg.tool_calls:",
        "        name = tc.function.name",
        "        args = json.loads(tc.function.arguments)",
        "        result = TOOL_FUNCTIONS[name](**args)",
        "        messages.append({",
        "            'role': 'tool', 'name': name,",
        "            'content': result, 'tool_call_id': tc.id,",
        "        })",
        "",
        "# ---------- actual step log ----------",
    ]
    for s in steps:
        lines.append(f"# Step {s['step']} [{s['type']}] tool={s['tool']}")
        lines.append(f"#   args:   {s['args']}")
        lines.append(f"#   result: {s['result']}")
    return "\n".join(lines)