# ============================================================================
# agent_workflow.py — Workflow backend (fixed 2-step prompt chain)
# ============================================================================
#
# CONTRACT (every backend file in this project exports these):
#   BACKEND_NAME: str
#   get_client(api_key: str) -> client
#   run(client, user_message: str) -> {"reply", "steps", "extracted"}
#   build_code_snippets(user_message: str, steps: list) -> str
#
# PATTERN
# -------
# Workflow is the simplest possible agentic structure: a fixed two-step
# prompt chain with NO tools. Step 1 clarifies the user's message. Step 2
# answers the clarified question. The developer, not the model, decides
# that there are exactly 2 steps in that exact order.
# ============================================================================

import os

# Defensive import: the mistralai package has been through THREE incompatible
# layouts and pip may install any of them depending on Python version and
# dependency resolution.
#   v2.x: from mistralai.client import Mistral         (latest, Nov 2025+)
#   v1.x: from mistralai import Mistral                (mid-2024 to late-2025)
#   v0.x: from mistralai.client import MistralClient   (pre-1.0)
# Try each in order and raise a clean error only if all three fail.
_Mistral = None
try:
    # v1.x: top-level import
    from mistralai import Mistral as _Mistral  # noqa: F401
except ImportError:
    try:
        # v2.x: moved to mistralai.client
        from mistralai.client import Mistral as _Mistral  # noqa: F401
    except ImportError:
        try:
            # v0.x: old class name in mistralai.client
            from mistralai.client import MistralClient as _OldClient
            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,
                    )

            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. "
                "Tried v1 (from mistralai import Mistral), "
                "v2 (from mistralai.client import Mistral), "
                "and v0 (from mistralai.client import MistralClient). "
                f"All failed. Last error: {_e}"
            )

Mistral = _Mistral

from parameters import TEMPERATURE, MAX_TOKENS
from prompts import WORKFLOW_STEP1_CLARIFY, WORKFLOW_STEP2_ANSWER
import providers


BACKEND_NAME = "Workflow"


def get_client(api_key, provider="Mistral"):
    """Return a provider-agnostic LLM client.

    The factory in providers.py handles all adapter logic. Old callers that
    pass only (api_key) still work — provider defaults to Mistral.
    """
    return providers.get_llm_client(provider, api_key)


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


def run(client, user_message, provider="Mistral"):
    """Fixed 2-step prompt chain: clarify -> answer. No tools."""
    steps = []

    step1 = _llm(client, [
        {"role": "system", "content": WORKFLOW_STEP1_CLARIFY},
        {"role": "user", "content": user_message},
    ], provider=provider)
    clarified = step1.content or ""
    steps.append({
        "step": 1, "type": "llm_call", "tool": "clarify",
        "args": user_message, "result": clarified,
    })

    step2 = _llm(client, [
        {"role": "system", "content": WORKFLOW_STEP2_ANSWER},
        {"role": "user", "content": clarified},
    ], provider=provider)
    answer = step2.content or ""
    steps.append({
        "step": 2, "type": "llm_call", "tool": "answer",
        "args": clarified, "result": answer,
    })

    return {
        "reply": answer,
        "steps": steps,
        "extracted": {"clarified_question": clarified},
    }


def build_code_snippets(user_message, steps):
    lines = [
        "# Backend: Workflow",
        "# Raw Mistral SDK, fixed 2-step prompt chain, no tools.",
        f"# User message: {user_message}",
        "",
        "# Step 1: clarify the user message using the clarify system prompt",
        "step1 = client.chat.complete(",
        "    model=MODEL,",
        "    messages=[",
        "        {'role': 'system', 'content': WORKFLOW_STEP1_CLARIFY},",
        f"        {{'role': 'user', 'content': {user_message!r}}},",
        "    ],",
        ").choices[0].message",
        "clarified = step1.content",
        "",
        "# Step 2: answer the clarified question using the answer system prompt",
        "step2 = client.chat.complete(",
        "    model=MODEL,",
        "    messages=[",
        "        {'role': 'system', 'content': WORKFLOW_STEP2_ANSWER},",
        "        {'role': 'user', 'content': clarified},",
        "    ],",
        ").choices[0].message",
        "answer = step2.content  # final reply to the user",
        "",
        "# ---------- actual step log ----------",
    ]
    for s in steps:
        lines.append(f"# Step {s['step']} [{s['type']}] {s['tool']}")
        lines.append(f"#   input:  {s['args']!r}")
        lines.append(f"#   output: {s['result']!r}")
    return "\n".join(lines)