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.routes.utils import not_found, parse_pagination
from fabledassistant.config import Config
from fabledassistant.services.chat import (
    add_message,
    bulk_delete_conversations,
    cleanup_old_conversations,
    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.notes import get_notes_by_ids
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, offset = parse_pagination()
    conv_type = request.args.get("type", "chat")
    # Apply retention policy before returning list
    retention_str = await get_setting(uid, "chat_retention_days", "90")
    try:
        retention_days = int(retention_str)
    except (ValueError, TypeError):
        retention_days = 90
    if retention_days > 0:
        await cleanup_old_conversations(uid, retention_days)
    conversations, total = await list_conversations(uid, limit=limit, offset=offset, conv_type=conv_type)
    return jsonify({
        "conversations": conversations,
        "total": total,
    })


@chat_bp.route("/conversations/bulk-delete", methods=["POST"])
@login_required
async def bulk_delete_conversations_route():
    uid = get_current_user_id()
    data = await request.get_json()
    ids = data.get("ids", []) if isinstance(data, dict) else []
    if not isinstance(ids, list) or not all(isinstance(i, int) for i in ids):
        return jsonify({"error": "ids must be a list of integers"}), 400
    count = await bulk_delete_conversations(uid, ids)
    return jsonify({"deleted": count})


@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)
    conversation_type = data.get("conversation_type", "chat")
    # Only allow known types to prevent accidental misuse
    if conversation_type not in ("chat", "mcp"):
        conversation_type = "chat"
    conv = await create_conversation(uid, title=title, model=model, conversation_type=conversation_type)
    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 not_found("Conversation")
    result = conv.to_dict()
    note_ids = [m.context_note_id for m in conv.messages if m.context_note_id]
    note_map = await get_notes_by_ids(uid, note_ids) if note_ids else {}
    result["messages"] = []
    for m in conv.messages:
        msg_dict = m.to_dict()
        if m.context_note_id and m.context_note_id in note_map:
            msg_dict["context_note_title"] = note_map[m.context_note_id].title
        result["messages"].append(msg_dict)
    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 not_found("Conversation")
    return "", 204


@chat_bp.route("/conversations/<int:conv_id>", methods=["PATCH"])
@login_required
async def update_conversation_route(conv_id: int):
    from fabledassistant.services.chat import _UNSET
    uid = get_current_user_id()
    data = await request.get_json()
    title = data.get("title")
    model = data.get("model")
    rag_project_id = data.get("rag_project_id", _UNSET)
    if title is None and model is None and rag_project_id is _UNSET:
        return jsonify({"error": "title, model, or rag_project_id is required"}), 400
    conv = await update_conversation(uid, conv_id, title=title, model=model, rag_project_id=rag_project_id)
    if conv is None:
        return not_found("Conversation")
    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 not_found("Conversation")

    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")
    include_note_ids = data.get("include_note_ids") or []
    excluded_note_ids = data.get("excluded_note_ids") or []
    think = bool(data.get("think", False))
    rag_project_id = data.get("rag_project_id") or None
    workspace_project_id = data.get("workspace_project_id") or None
    user_timezone = data.get("user_timezone") or None
    if not user_timezone:
        user_timezone = await get_setting(uid, "user_timezone") or None
    effective_rag_project_id = workspace_project_id or rag_project_id

    # 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")

    try:
        buf = create_buffer(conv_id, assistant_msg.id)
    except RuntimeError:
        return jsonify({"error": "Generation already in progress"}), 409

    # 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 = await get_setting(uid, "default_model", Config.OLLAMA_MODEL) or 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,
        include_note_ids=include_note_ids,
        excluded_note_ids=excluded_note_ids,
        think=think,
        rag_project_id=effective_rag_project_id,
        workspace_project_id=workspace_project_id,
        user_timezone=user_timezone,
    ))

    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 not_found("Conversation")

    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")
    try:
        last_id = int(last_id_str) if last_id_str is not None else -1
    except (ValueError, TypeError):
        last_id = -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"

        # Final drain: deliver any events appended between the last yield
        # and the state check above (e.g. the 'done' event).
        for event in buf.events_after(cursor):
            yield buf.format_sse(event)

    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 not_found("Conversation")

    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 not_found("Conversation")

    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 not_found("Conversation")
    model = await get_setting(uid, "default_model", Config.OLLAMA_MODEL) or 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)
    # Guard against empty-string rows written by older code when user selected "default"
    if not default_model:
        default_model = Config.OLLAMA_MODEL
    result = {
        "ollama": "unavailable",
        "model": "not_found",
        "default_model": default_model,
    }
    try:
        async with httpx.AsyncClient(timeout=10.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 or base in model_names:
                    # Installed — now check if currently loaded in memory
                    result["model"] = "cold"
                    if not isinstance(ps_resp, Exception):
                        try:
                            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 or base in loaded_names:
                                result["model"] = "loaded"
                        except Exception:
                            logger.debug("Ollama /api/ps check failed", exc_info=True)
    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:
            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)

        loaded_names: set[str] = set()
        if not isinstance(ps_resp, Exception):
            try:
                ps_resp.raise_for_status()
                loaded_names = {m["name"] for m in ps_resp.json().get("models", [])}
            except Exception:
                pass

        models = []
        if not isinstance(tags_resp, Exception):
            tags_resp.raise_for_status()
            for m in tags_resp.json().get("models", []):
                models.append({
                    "name": m["name"],
                    "size": m.get("size", 0),
                    "modified_at": m.get("modified_at", ""),
                    "loaded": m["name"] in loaded_names,
                })
        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