# app/services/ollama.py
import httpx
import json
from app.core.config import settings
from app.core.context import SYSTEM_PROMPT

_histories: dict[str, list[dict]] = {}


def reset_session(session_id: str) -> None:
    _histories.pop(session_id, None)


async def chat(user_message: str, session_id: str) -> tuple[str, str]:
    """
    Fast Ollama call with aggressive token limits and keep_alive caching.
    Returns (reply, source).
    """
    history = _histories.setdefault(session_id, [])
    max_msgs = settings.MAX_HISTORY * 2
    if len(history) > max_msgs:
        history = history[-max_msgs:]

    messages = (
        [{"role": "system", "content": SYSTEM_PROMPT}]
        + history
        + [{"role": "user", "content": user_message}]
    )

    payload = {
        "model":   settings.OLLAMA_MODEL,
        "messages": messages,
        "stream":  False,
        "keep_alive": "10m",        # keep model loaded in RAM between requests
        "options": {
            "temperature":  0.1,
            "num_predict":  1024,
            "num_ctx":      4096,
            "top_k":        10, 
            "repeat_penalty": 1.1,
    }
    }

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            resp = await client.post(
                f"{settings.OLLAMA_URL}/api/chat",
                json=payload,
            )
            resp.raise_for_status()
            reply = resp.json()["message"]["content"].strip()

        history.append({"role": "user",      "content": user_message})
        history.append({"role": "assistant", "content": reply})
        _histories[session_id] = history
        return reply, "llm"

    except httpx.ConnectError:
        return "⚠️ Ollama is not running. Run: ollama serve", "error"
    except httpx.TimeoutException:
        return "Response timed out — try a smaller model like phi3:mini.", "error"
    except Exception as e:
        return f"Error: {e}", "error"