import json
from dataclasses import asdict, dataclass
from urllib import request, error

from aiohttp import web

from server import PromptServer

OLLAMA_HOST = "http://127.0.0.1:11434"
OLLAMA_URL_GENERATE = f"{OLLAMA_HOST}/api/generate"
OLLAMA_URL_CHAT = f"{OLLAMA_HOST}/api/chat"
OLLAMA_TAGS_URL = f"{OLLAMA_HOST}/api/tags"


# Server api definitions

@PromptServer.instance.routes.get("/prompt_bench/ollama/models")
async def get_ollama_models(_request):
    return web.json_response(_fetch_ollama_models())


# Custom request definitions

def _fetch_ollama_models() -> list[str]:
    print("fetching ollama models")
    try:
        with request.urlopen(OLLAMA_TAGS_URL, timeout=5) as response:
            print("received ollama models")
            body = json.loads(response.read().decode("utf-8"))
            models = [
                model.get("name", "").strip()
                for model in body.get("models", [])
                if model.get("name", "").strip()
            ]
            return models
    except Exception:
        print("fetching ollama models failed")
        return []


def post_ollama_text(model: str, prompt: str):
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
    }

    req = request.Request(
        OLLAMA_URL_GENERATE,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    try:
        with request.urlopen(req, timeout=120) as response:
            body = json.loads(response.read().decode("utf-8"))
    except error.URLError as exc:
        print(exc)
        # Ignore error

    if "error" in body:
        raise RuntimeError(f"Ollama generation failed: {body['error']}")

    text = body.get("response", "").strip()
    if not text:
        raise RuntimeError("Ollama returned an empty response.")

    return text


@dataclass
class OllamaChatMessage:
    role: str
    content: str


def post_ollama_chat(model: str, messages: list[OllamaChatMessage]):
    payload = {
        "model": model,
        "messages": [asdict(m) for m in messages],
        "stream": False,
    }

    req = request.Request(
        OLLAMA_URL_CHAT,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    try:
        with request.urlopen(req, timeout=120) as response:
            body = json.loads(response.read().decode("utf-8"))
    except error.URLError as exc:
        print(exc)
        # Ignore error

    if "error" in body:
        raise RuntimeError(f"Ollama generation failed: {body['error']}")

    assistant_message_raw = body.get("message", {"role": "assistant", "content": None})
    if not assistant_message_raw.get("content", None):
        raise RuntimeError("Ollama returned an empty response.")

    return OllamaChatMessage(role="assistant", content=assistant_message_raw["content"])