#!/usr/bin/env python
"""LLM backend for the DIY Creator multi-agent team (#17): local Ollama (default) or Claude.

Two entry points:
  chat(system, user)       -> free text (the writer / reviser produce markdown guides)
  json_call(system, user)  -> parsed JSON (the researcher's brief, the critic's verdict)

Each call takes an optional `model` so different ROLES can run on different models — which is the
whole point of the critic: it can be a *different* model from the writer, so the review is
independent (the "diamond"), not the author grading its own work. Errors log the exception TYPE only.
"""
import json
import os
import urllib.request

BACKEND      = os.environ.get("DIY_LLM_BACKEND", "ollama").lower()
OLLAMA_URL   = os.environ.get("DIY_OLLAMA_URL", "http://127.0.0.1:11434").rstrip("/")
OLLAMA_MODEL = os.environ.get("DIY_OLLAMA_MODEL", "qwen2.5:7b")
# The critic defaults to the same model, but set DIY_CRITIC_MODEL to a *different* one for a truly
# independent review (e.g. a different family, or a bigger model where the quality bar matters).
CRITIC_MODEL = os.environ.get("DIY_CRITIC_MODEL", "").strip() or OLLAMA_MODEL

ANTHROPIC_KEY    = (os.environ.get("ANTHROPIC_API_KEY") or "").strip() or None
ANTHROPIC_MODEL  = os.environ.get("ANTHROPIC_MODEL", "claude-haiku-4-5")
ANTHROPIC_CRITIC = os.environ.get("ANTHROPIC_CRITIC_MODEL", "").strip() or ANTHROPIC_MODEL

_USE_CLAUDE = BACKEND == "claude" and ANTHROPIC_KEY


def _ollama(system, user, model, temperature, max_tokens, as_json):
    payload = {"model": model, "stream": False,
               "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}],
               "options": {"temperature": temperature, "num_predict": max_tokens}}
    if as_json:
        payload["format"] = "json"          # Ollama constrains output to valid JSON
    body = json.dumps(payload).encode()
    req = urllib.request.Request(OLLAMA_URL + "/api/chat", data=body,
                                 headers={"Content-Type": "application/json"})
    return json.loads(urllib.request.urlopen(req, timeout=180).read())["message"]["content"].strip()


def _claude(system, user, model, temperature, max_tokens, as_json):
    if as_json:
        system = system + "\nReply with ONLY valid JSON — no prose, no code fences."
    body = json.dumps({"model": model, "max_tokens": max_tokens, "temperature": temperature,
                       "system": system, "messages": [{"role": "user", "content": user}]}).encode()
    req = urllib.request.Request("https://api.anthropic.com/v1/messages", data=body, headers={
        "x-api-key": ANTHROPIC_KEY, "anthropic-version": "2023-06-01", "content-type": "application/json"})
    data = json.loads(urllib.request.urlopen(req, timeout=180).read())
    return "".join(b.get("text", "") for b in data.get("content", [])).strip()


def _resolve(model, role):
    """Pick the concrete model id for a role given the active backend."""
    if _USE_CLAUDE:
        return ANTHROPIC_CRITIC if role == "critic" else ANTHROPIC_MODEL
    if role == "critic":
        return CRITIC_MODEL
    return model or OLLAMA_MODEL


def chat(system: str, user: str, role: str = "", model: str = None,
         temperature: float = 0.5, max_tokens: int = 1400) -> str:
    m = _resolve(model, role)
    fn = _claude if _USE_CLAUDE else _ollama
    return fn(system, user, m, temperature, max_tokens, as_json=False)


def json_call(system: str, user: str, role: str = "", model: str = None,
              temperature: float = 0.2, max_tokens: int = 1200):
    m = _resolve(model, role)
    fn = _claude if _USE_CLAUDE else _ollama
    raw = fn(system, user, m, temperature, max_tokens, as_json=True)
    return _parse_json(raw)


def _parse_json(raw: str):
    s = (raw or "").strip()
    if s.startswith("```"):
        s = s.split("```", 2)[1].lstrip("json").strip() if s.count("```") >= 2 else s.strip("`")
    try:
        return json.loads(s)
    except Exception:
        pass
    # salvage the first {...} or [...] block
    for open_c, close_c in (("{", "}"), ("[", "]")):
        i, j = s.find(open_c), s.rfind(close_c)
        if 0 <= i < j:
            try:
                return json.loads(s[i:j + 1])
            except Exception:
                continue
    return None


def backend_label() -> str:
    if _USE_CLAUDE:
        extra = f" · critic: {ANTHROPIC_CRITIC}" if ANTHROPIC_CRITIC != ANTHROPIC_MODEL else ""
        return f"Claude ({ANTHROPIC_MODEL}){extra}"
    extra = f" · critic: {CRITIC_MODEL}" if CRITIC_MODEL != OLLAMA_MODEL else ""
    return f"Ollama ({OLLAMA_MODEL} @ {OLLAMA_URL}){extra}"