"""
Ollama local API — call Ollama to generate replies, parsed as ASDP action_queue
Returns (content, usage) where usage has prompt_tokens, completion_tokens, total_tokens.
"""
from __future__ import annotations

import json
import urllib.request
import urllib.error
from typing import Any

DEFAULT_BASE = "http://localhost:11434"


def _usage_from_ollama(out: dict[str, Any]) -> dict[str, int]:
    """Extract token usage from Ollama /api/chat response."""
    prompt = out.get("prompt_eval_count") or 0
    completion = out.get("eval_count") or 0
    return {"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion}


def call_cloud_api(
    model: str,
    system_prompt: str,
    user_message: str,
    api_key: str,
    base_url: str = "https://api.siliconflow.cn/v1"
) -> tuple[str, dict[str, int]]:
    url = f"{base_url.rstrip('/')}/chat/completions"
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
        "stream": False,
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"  # cloud API requires key
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            out = json.loads(resp.read().decode())
            content = out.get("choices", [{}])[0].get("message", {}).get("content", "")
            usage = out.get("usage") or {}
            u = {
                "prompt_tokens": usage.get("prompt_tokens", 0),
                "completion_tokens": usage.get("completion_tokens", 0),
                "total_tokens": usage.get("total_tokens", 0),
            }
            return content, u
    except Exception as e:
        print(f"  [API error] {e}")
        return '{"action_queue": [{"type": "idle", "parameters": null}]}', {}


def call_ollama(
    model: str,
    system_prompt: str,
    user_message: str,
    base_url: str = DEFAULT_BASE,
) -> tuple[str, dict[str, int]]:
    """
    Call Ollama /api/chat, return (assistant content, usage).
    On failure, return (fallback JSON, {}).
    """
    url = f"{base_url.rstrip('/')}/api/chat"
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
        "stream": False,
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            out = json.loads(resp.read().decode())
            content = out.get("message", {}).get("content", "") or ""
            usage = _usage_from_ollama(out)
            return content, usage
    except (urllib.error.URLError, OSError, json.JSONDecodeError) as e:
        print(f"  [Ollama error] {e}")
        return '{"action_queue": [{"type": "idle", "parameters": null}]}', {}
