import asyncio
from typing import Any, Dict, List, Optional
from django.conf import settings
from mcp_agent.mcp_client import MCPClient
from .models import AgentModel


async def _call_mcp(tool: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Internal async helper to call the MCP HTTP bridge via MCPClient."""
    server_url = getattr(settings, "MCP_AGENT_URL")
    client = MCPClient(server_url)
    try:
        resp = await client.send(tool, arguments)
        return resp
    finally:
        await client.close()


def fine_tune_model(
    base_model: str,
    training_files: List[str],
    hyperparams: Dict[str, Any],
    name: str,
    version: str,
) -> Dict[str, Any]:
    """Synchronously request a fine-tune run on the MCP server.

    Expects the MCP tool `fine_tune` to accept: {base_model, training_files, hyperparams, name, version}
    and to return a JSON-like dict containing at least `status` and on success `model_path` and `version`.
    """
    return asyncio.run(_call_mcp("fine_tune", {
        "base_model": base_model,
        "training_files": training_files,
        "hyperparams": hyperparams,
        "name": name,
        "version": version,
    }))


def load_model_for_inference(model_path: str) -> Dict[str, Any]:
    """Tell the MCP server to load a model into memory/serving for inference.

    Expects the MCP tool `load_model` with {model_path} returning status info.
    """
    return asyncio.run(_call_mcp("load_model", {"model_path": model_path}))


def infer_with_model(model_path: str, prompt: str, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Request inference from the MCP server using a previously fine-tuned model.

    Calls the MCP tool `infer` with {model_path, prompt, options}.
    """
    return asyncio.run(_call_mcp("infer", {"model_path": model_path, "prompt": prompt, "options": options or {}}))


def register_model_in_db(name: str, version: str, model_path: str) -> AgentModel:
    """Convenience DB helper: create and return an AgentModel record.

    NOTE: migrations are required after the model field change prior to using this in production.
    """
    return AgentModel.objects.create(name=name, version=version, path=model_path)