FabledScribe/src/fabledassistant/routes/chat.py

import asyncio
import json
import logging

import httpx
from quart import Blueprint, Response, jsonify, request

from fabledassistant.auth import login_required, get_current_user_id
from fabledassistant.config import Config
from fabledassistant.services.chat import (
    add_message,
    create_conversation,
    delete_conversation,
    get_conversation,
    list_conversations,
    save_response_as_note,
    summarize_conversation_as_note,
    update_conversation,
)
from fabledassistant.services.generation_buffer import (
    GenerationState,
    create_buffer,
    get_buffer,
)
from fabledassistant.services.generation_task import run_generation
from fabledassistant.services.settings import get_setting

logger = logging.getLogger(__name__)

chat_bp = Blueprint("chat", __name__, url_prefix="/api/chat")


@chat_bp.route("/conversations", methods=["GET"])
@login_required
async def list_conversations_route():
    uid = get_current_user_id()
    limit = request.args.get("limit", 50, type=int)
    offset = request.args.get("offset", 0, type=int)
    conversations, total = await list_conversations(uid, limit=limit, offset=offset)
    return jsonify({
        "conversations": conversations,
        "total": total,
    })


@chat_bp.route("/conversations", methods=["POST"])
@login_required
async def create_conversation_route():
    uid = get_current_user_id()
    data = await request.get_json(force=True, silent=True) or {}
    title = data.get("title", "")
    model = data.get("model", Config.OLLAMA_MODEL)
    conv = await create_conversation(uid, title=title, model=model)
    return jsonify(conv.to_dict()), 201


@chat_bp.route("/conversations/<int:conv_id>", methods=["GET"])
@login_required
async def get_conversation_route(conv_id: int):
    uid = get_current_user_id()
    conv = await get_conversation(uid, conv_id)
    if conv is None:
        return jsonify({"error": "Conversation not found"}), 404
    result = conv.to_dict()
    result["messages"] = [m.to_dict() for m in conv.messages]
    return jsonify(result)


@chat_bp.route("/conversations/<int:conv_id>", methods=["DELETE"])
@login_required
async def delete_conversation_route(conv_id: int):
    uid = get_current_user_id()
    deleted = await delete_conversation(uid, conv_id)
    if not deleted:
        return jsonify({"error": "Conversation not found"}), 404
    return "", 204


@chat_bp.route("/conversations/<int:conv_id>", methods=["PATCH"])
@login_required
async def update_conversation_route(conv_id: int):
    uid = get_current_user_id()
    data = await request.get_json()
    title = data.get("title")
    model = data.get("model")
    if title is None and model is None:
        return jsonify({"error": "title or model is required"}), 400
    conv = await update_conversation(uid, conv_id, title=title, model=model)
    if conv is None:
        return jsonify({"error": "Conversation not found"}), 404
    return jsonify(conv.to_dict())


@chat_bp.route("/conversations/<int:conv_id>/messages", methods=["POST"])
@login_required
async def send_message_route(conv_id: int):
    """Start generation: save user message, launch background task, return 202."""
    uid = get_current_user_id()
    conv = await get_conversation(uid, conv_id)
    if conv is None:
        return jsonify({"error": "Conversation not found"}), 404

    data = await request.get_json()
    content = data.get("content", "").strip()
    if not content:
        return jsonify({"error": "content is required"}), 400
    context_note_id = data.get("context_note_id")
    exclude_note_ids = data.get("exclude_note_ids") or []

    # Reject if generation already running for this conversation
    existing = get_buffer(conv_id)
    if existing and existing.state == GenerationState.RUNNING:
        return jsonify({"error": "Generation already in progress"}), 409

    # Save user message
    await add_message(conv_id, "user", content, context_note_id=context_note_id)

    # Create placeholder assistant message
    assistant_msg = await add_message(conv_id, "assistant", "", status="generating")

    buf = create_buffer(conv_id, assistant_msg.id)

    # Build history from existing messages (excluding system and the placeholder)
    history = []
    for msg in conv.messages:
        if msg.role != "system":
            history.append({"role": msg.role, "content": msg.content})

    model = conv.model or await get_setting(uid, "default_model", Config.OLLAMA_MODEL)

    # Launch background generation task (context building happens inside the task)
    asyncio.create_task(run_generation(
        buf, history, model,
        uid, conv_id, conv.title, content,
        context_note_id=context_note_id,
        exclude_note_ids=exclude_note_ids,
    ))

    return jsonify({
        "assistant_message_id": assistant_msg.id,
        "status": "generating",
    }), 202


@chat_bp.route("/conversations/<int:conv_id>/generation/stream", methods=["GET"])
@login_required
async def generation_stream_route(conv_id: int):
    """SSE endpoint that tails the generation buffer for a conversation."""
    uid = get_current_user_id()
    conv = await get_conversation(uid, conv_id)
    if conv is None:
        return jsonify({"error": "Conversation not found"}), 404

    buf = get_buffer(conv_id)
    if buf is None:
        return jsonify({"error": "No active generation"}), 404

    # Determine starting point from Last-Event-ID header or query param
    last_id_str = request.headers.get("Last-Event-ID") or request.args.get("last_event_id")
    last_id = int(last_id_str) if last_id_str is not None else -1

    async def stream():
        cursor = last_id
        while True:
            # Replay any buffered events past the cursor
            pending = buf.events_after(cursor)
            for event in pending:
                yield buf.format_sse(event)
                cursor = event.index

            # If generation is done and all events delivered, close stream
            if buf.state != GenerationState.RUNNING:
                break

            # Wait for new events or send keepalive on timeout
            got_new = await buf.wait_for_event(cursor, timeout=15.0)
            if not got_new:
                yield ": keepalive\n\n"

    return Response(
        stream(),
        content_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


@chat_bp.route("/conversations/<int:conv_id>/generation/confirm", methods=["POST"])
@login_required
async def confirm_generation_route(conv_id: int):
    """Resolve a pending tool confirmation (accept or decline)."""
    uid = get_current_user_id()
    conv = await get_conversation(uid, conv_id)
    if conv is None:
        return jsonify({"error": "Conversation not found"}), 404

    buf = get_buffer(conv_id)
    if buf is None or buf.state != GenerationState.RUNNING:
        return jsonify({"error": "No active generation"}), 404

    if buf.confirmation_future is None or buf.confirmation_future.done():
        return jsonify({"error": "No pending tool confirmation"}), 409

    data = await request.get_json(force=True, silent=True) or {}
    decision = bool(data.get("confirmed", False))
    buf.confirmation_future.set_result(decision)
    return jsonify({"status": "ok", "confirmed": decision})


@chat_bp.route("/conversations/<int:conv_id>/generation/cancel", methods=["POST"])
@login_required
async def cancel_generation_route(conv_id: int):
    """Cancel an active generation for a conversation."""
    uid = get_current_user_id()
    conv = await get_conversation(uid, conv_id)
    if conv is None:
        return jsonify({"error": "Conversation not found"}), 404

    buf = get_buffer(conv_id)
    if buf is None or buf.state != GenerationState.RUNNING:
        return jsonify({"error": "No active generation"}), 404

    buf.cancel_event.set()
    return jsonify({"status": "cancelled"})


@chat_bp.route("/messages/<int:message_id>/save-as-note", methods=["POST"])
@login_required
async def save_message_as_note_route(message_id: int):
    uid = get_current_user_id()
    try:
        note = await save_response_as_note(uid, message_id)
        return jsonify(note), 201
    except ValueError as e:
        return jsonify({"error": str(e)}), 400


@chat_bp.route("/conversations/<int:conv_id>/summarize", methods=["POST"])
@login_required
async def summarize_conversation_route(conv_id: int):
    uid = get_current_user_id()
    conv = await get_conversation(uid, conv_id)
    if conv is None:
        return jsonify({"error": "Conversation not found"}), 404
    model = conv.model or await get_setting(uid, "default_model", Config.OLLAMA_MODEL)
    try:
        note = await summarize_conversation_as_note(uid, conv_id, model)
        return jsonify(note), 201
    except ValueError as e:
        return jsonify({"error": str(e)}), 400


@chat_bp.route("/ps", methods=["GET"])
@login_required
async def running_models_route():
    """Return currently loaded (hot) models from Ollama."""
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            resp = await client.get(f"{Config.OLLAMA_URL}/api/ps")
            resp.raise_for_status()
            data = resp.json()
            models = [
                {
                    "name": m["name"],
                    "size": m.get("size", 0),
                    "size_vram": m.get("size_vram", 0),
                    "expires_at": m.get("expires_at", ""),
                }
                for m in data.get("models", [])
            ]
            return jsonify({"models": models})
    except Exception as e:
        logger.debug("Failed to query Ollama /api/ps: %s", e)
        return jsonify({"models": []})


@chat_bp.route("/warm", methods=["POST"])
@login_required
async def warm_model_route():
    """Pre-load a model into Ollama memory."""
    data = await request.get_json(force=True, silent=True) or {}
    model = data.get("model", "").strip()
    if not model:
        return jsonify({"error": "model is required"}), 400

    async def _warm():
        try:
            async with httpx.AsyncClient(timeout=300.0) as client:
                await client.post(
                    f"{Config.OLLAMA_URL}/api/generate",
                    json={"model": model, "prompt": "", "keep_alive": "30m"},
                )
            logger.info("Warmed model %s", model)
        except Exception as e:
            logger.warning("Failed to warm model %s: %s", model, e)

    asyncio.create_task(_warm())
    return jsonify({"status": "warming"}), 202


@chat_bp.route("/status", methods=["GET"])
@login_required
async def chat_status_route():
    """Check Ollama availability, model installation, and model load state."""
    uid = get_current_user_id()
    default_model = await get_setting(uid, "default_model", Config.OLLAMA_MODEL)
    result = {
        "ollama": "unavailable",
        "model": "not_found",
        "default_model": default_model,
    }
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            # Query installed models and loaded models in parallel
            tags_task = asyncio.create_task(client.get(f"{Config.OLLAMA_URL}/api/tags"))
            ps_task = asyncio.create_task(client.get(f"{Config.OLLAMA_URL}/api/ps"))
            tags_resp, ps_resp = await asyncio.gather(tags_task, ps_task, return_exceptions=True)

            if isinstance(tags_resp, Exception):
                logger.debug("Ollama /api/tags failed: %s", tags_resp)
            else:
                tags_resp.raise_for_status()
                result["ollama"] = "available"
                model_names = {m["name"] for m in tags_resp.json().get("models", [])}
                base = default_model.removesuffix(":latest")
                if default_model in model_names or f"{base}:latest" in model_names:
                    # Installed — now check if currently loaded in memory
                    result["model"] = "cold"
                    if not isinstance(ps_resp, Exception):
                        ps_resp.raise_for_status()
                        loaded_names = {m["name"] for m in ps_resp.json().get("models", [])}
                        if default_model in loaded_names or f"{base}:latest" in loaded_names:
                            result["model"] = "loaded"
    except Exception:
        logger.debug("Ollama status check failed", exc_info=True)
    return jsonify(result)


@chat_bp.route("/models", methods=["GET"])
@login_required
async def list_models_route():
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.get(f"{Config.OLLAMA_URL}/api/tags")
            resp.raise_for_status()
            data = resp.json()
            models = [
                {"name": m["name"], "size": m.get("size", 0)}
                for m in data.get("models", [])
            ]
            return jsonify({"models": models})
    except Exception as e:
        logger.warning("Failed to list Ollama models: %s", e)
        return jsonify({"models": [], "error": str(e)}), 200


@chat_bp.route("/models/pull", methods=["POST"])
@login_required
async def pull_model_route():
    """Pull a model from Ollama, streaming progress via SSE."""
    data = await request.get_json()
    model_name = data.get("model", "").strip()
    if not model_name:
        return jsonify({"error": "model is required"}), 400

    async def generate():
        try:
            async with httpx.AsyncClient(timeout=1800.0) as client:
                async with client.stream(
                    "POST",
                    f"{Config.OLLAMA_URL}/api/pull",
                    json={"name": model_name},
                ) as resp:
                    resp.raise_for_status()
                    async for line in resp.aiter_lines():
                        if not line.strip():
                            continue
                        progress = json.loads(line)
                        event_data = json.dumps(progress)
                        yield f"data: {event_data}\n\n"
            yield f"data: {json.dumps({'status': 'success'})}\n\n"
        except Exception as e:
            logger.exception("Error pulling model %s", model_name)
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    return Response(
        generate(),
        content_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


@chat_bp.route("/models/delete", methods=["POST"])
@login_required
async def delete_model_route():
    """Delete a model from Ollama."""
    data = await request.get_json()
    model_name = data.get("model", "").strip()
    if not model_name:
        return jsonify({"error": "model is required"}), 400

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.request(
                "DELETE",
                f"{Config.OLLAMA_URL}/api/delete",
                json={"name": model_name},
            )
            resp.raise_for_status()
        return jsonify({"status": "deleted", "model": model_name})
    except httpx.HTTPStatusError as e:
        logger.warning("Failed to delete model %s: %s", model_name, e)
        return jsonify({"error": f"Failed to delete model: {e.response.status_code}"}), 400
    except Exception as e:
        logger.warning("Failed to delete model %s: %s", model_name, e)
        return jsonify({"error": str(e)}), 500