from axon.backends.base_backend import BaseBackend, CompiledStep, CompilationContext
from axon.compiler.ir_nodes import IRNode, IRPersona, IRContext, IRAnchor, IRToolSpec

class OllamaBackend(BaseBackend):
    """Stub implementation for the Ollama backend."""

    @property
    def name(self) -> str:
        return "ollama"

    def compile_step(
        self, step: IRNode, context: CompilationContext
    ) -> CompiledStep:
        raise NotImplementedError(
            "Ollama backend is not yet implemented. "
            "Scheduled for Phase 2 expansion. "
            "Should adapt prompts for local models with smaller "
            "context windows and optional tool support."
        )

    def compile_system_prompt(
        self,
        persona: IRPersona | None,
        context: IRContext | None,
        anchors: list[IRAnchor],
    ) -> str:
        raise NotImplementedError(
            "Ollama system prompt compilation is not yet implemented. "
            "Should produce simplified system prompts suitable for "
            "local models (Llama, Mistral, etc.)."
        )

    def compile_tool_spec(self, tool: IRToolSpec) -> dict[str, Any]:
        raise NotImplementedError(
            "Ollama tool spec compilation is not yet implemented. "
            "Should produce Ollama-compatible tool format or gracefully "
            "degrade for models without tool support."
        )

[INST] <<SYS>>
You are a helpful assistant.
<</SYS>>

User question here [/INST]

<s>[INST] User instruction [/INST]

### Instruction:
You are a helpful assistant.

### Input:
User question here

### Response:

def compile_system_prompt(
    self, persona, context, anchors
) -> str:
    # Simplified format for local models
    parts = []
    
    if persona:
        # Concise persona: "You are X specializing in Y."
        parts.append(f"You are {persona.name} specializing in {', '.join(persona.domain[:2])}.")
    
    if anchors:
        # Compact constraint list
        parts.append("Rules: " + "; ".join(a.require for a in anchors if a.require))
    
    return " ".join(parts)

You are LegalExpert specializing in contract law, IP. Rules: cite all sources; no hallucination.

def _detect_instruction_format(self, model_name: str) -> str:
    """Detect instruction template from model name."""
    if "llama" in model_name.lower():
        return "llama2_chat"
    if "mistral" in model_name.lower():
        return "mistral_instruct"
    if "alpaca" in model_name.lower():
        return "alpaca"
    return "generic"

def _apply_instruction_template(
    self, prompt: str, model_format: str
) -> str:
    """Wrap prompt in model-specific instruction template."""
    if model_format == "llama2_chat":
        return f"[INST] {prompt} [/INST]"
    elif model_format == "mistral_instruct":
        return f"<s>[INST] {prompt} [/INST]"
    # ...
    return prompt

def compile_tool_spec(self, tool: IRToolSpec) -> dict[str, Any]:
    # Check if model supports tools (hypothetical API)
    if self._model_supports_tools():
        # Return Ollama tool format
        return {"name": tool.name, "description": tool.description}
    else:
        # Degrade to text-based tool calling
        return {
            "name": tool.name,
            "mode": "text_simulation",
            "instruction": (
                f"To use {tool.name}, output: TOOL_CALL[{tool.name}](query=\"...\")\n"
                f"The system will execute it and provide results."
            ),
        }

import httpx

class OllamaModelClient:
    """Model client for Ollama API."""
    
    def __init__(self, base_url: str = "http://localhost:11434"):
        self.base_url = base_url
    
    async def call(
        self,
        system_prompt: str,
        user_prompt: str,
        model: str = "llama2",
        **kwargs
    ) -> ModelResponse:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.base_url}/api/generate",
                json={
                    "model": model,
                    "prompt": self._format_prompt(system_prompt, user_prompt),
                    "stream": False,
                },
            )
            data = response.json()
            return ModelResponse(content=data["response"])
    
    def _format_prompt(self, system: str, user: str) -> str:
        # Apply Llama 2 instruction template
        return f"[INST] <<SYS>>\n{system}\n<</SYS>>\n\n{user} [/INST]"