import requests
import json
import re

LOCAL_GPU_MODEL = "gpt-oss:120b"

def ask_ollama(prompt: str, model: str = "qwen2.5-coder:1.5b", keep_alive=None) -> str:
    """
    呼叫本地 Ollama，回傳文字（不串流）
    """
    url = "http://localhost:11434/api/generate"
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False
    }
    if keep_alive is not None:
        payload["keep_alive"] = keep_alive
    resp = requests.post(url, json=payload)
    resp.raise_for_status()
    return resp.json().get("response", "").strip()

def unload_ollama_model(model: str) -> bool:
    """Unload only the local 120b model; cloud model names are intentionally ignored."""
    if model != LOCAL_GPU_MODEL:
        return False
    url = "http://localhost:11434/api/generate"
    payload = {
        "model": model,
        "prompt": "",
        "stream": False,
        "keep_alive": 0,
    }
    resp = requests.post(url, json=payload, timeout=(5, 60))
    resp.raise_for_status()
    return True

'''
def ask_ollama_docker(prompt: str, model: str = "openai/gpt-oss-20b") -> str:
    """
    呼叫 Docker 部署的 gpt-oss (OpenAI-compatible API)，回傳文字（不串流）
    """
    url = "http://192.168.41.30:8520/v1/chat/completions"
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.8,
        "max_tokens": 2048,
        "stream": False
    }

    resp = requests.post(
        url,
        headers={"Content-Type": "application/json"},
        json=payload,
        timeout=(10, 600),
    )
    resp.raise_for_status()

    data = resp.json()

    try:
        msg = (data.get("choices") or [{}])[0].get("message") or {}

        print(msg)

        # 1) 先拿 content；沒有就拿 reasoning_content；再不行才退回 text
        raw = (
            msg.get("content")
            or msg.get("reasoning_content")
            or (data.get("choices") or [{}])[0].get("text")
            or ""
        )
        raw = str(raw).strip()

        # 2) 只保留 MAP: 行，確保 parse_llm_maps() 一定吃得到
        map_lines = [ln.strip() for ln in raw.splitlines() if ln.strip().startswith("MAP:")]
        return "\n".join(map_lines).strip()

    except Exception:
        raise RuntimeError(f"Unexpected response format: {data}")
'''
def ask_ollama_docker(prompt: str, model: str = "openai/gpt-oss-20b") -> str:
    """
    呼叫 Docker 部署的 OpenAI-compatible API（/v1/completions）
    ✅ 回傳只保留 MAP: 行，避免 We need... 這種分析文字干擾下游
    """
    import requests

    url = "http://192.168.41.30:8520/v1/completions"
    payload = {
        "model": model,
        "prompt": prompt,
        "temperature": 0.2,
        "max_tokens": 512,
        "stream": False,
        # （可選）有些服務支援 stop，能減少碎念；不保證每家都有效
        # "stop": ["\nWe need", "\nWe ", "\nAnalysis", "\n\nWe"],
    }

    resp = requests.post(
        url,
        headers={"Content-Type": "application/json"},
        json=payload,
        timeout=(10, 600),
    )
    resp.raise_for_status()
    data = resp.json()

    raw = (data.get("choices") or [{}])[0].get("text") or ""
    raw = str(raw)

    # ✅ 只取 MAP: 行
    map_lines = []
    for ln in raw.splitlines():
        ln = ln.strip()
        if ln.startswith("MAP:"):
            map_lines.append(ln)

    return "\n".join(map_lines).strip()

def _extract_first_json(text: str) -> dict:
    """
    從模型輸出中抓第一個 JSON object（容錯用）
    """
    if not text:
        raise ValueError("empty llm output")

    # 最常見：模型前後夾雜文字，抓第一個 {...}
    m = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if not m:
        raise ValueError(f"no json found in llm output: {text[:200]}")
    return json.loads(m.group(0))

def ask_ollama_json(prompt: str, model: str = "gpt-oss:120b-cloud") -> dict:
    """
    用本機 Ollama 當 planner：強制輸出 JSON
    """
    # 讓模型更願意輸出純 JSON（比你自己 parse 更穩）
    wrapped = (
        "你只能輸出 JSON，禁止輸出任何其他文字。\n"
        "JSON 必須可被 json.loads 解析。\n\n"
        + prompt
    )
    raw = ask_ollama(wrapped, model=model)   # 你原本的 ask_ollama()
    return _extract_first_json(raw)