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 admin_required, 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,
    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")
    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", "voice"):
        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":
            continue
        msg_dict = {"role": msg.role, "content": msg.content or ""}
        if msg.tool_calls:
            msg_dict["tool_calls"] = [
                {"function": {"name": tc["function"], "arguments": tc["arguments"]}}
                for tc in msg.tool_calls
            ]
            history.append(msg_dict)
            for tc in msg.tool_calls:
                history.append({"role": "tool", "content": json.dumps(tc.get("result", {}))})
        else:
            history.append(msg_dict)

    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,
        voice_mode=(conv.conversation_type == "voice"),
    ))

    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():
        from fabledassistant.services.llm import keep_alive_for
        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": keep_alive_for(model)},
                )
            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"])
@admin_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"])
@admin_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


@chat_bp.route("/from-article/<int:item_id>", methods=["POST"])
@login_required
async def create_conversation_from_article(item_id: int):
    """Create a chat conversation seeded with an RSS article's content."""
    from sqlalchemy import select as _select
    from fabledassistant.models import async_session as _async_session
    from fabledassistant.models.conversation import Message
    from fabledassistant.models.rss_feed import RssItem, RssFeed

    uid = get_current_user_id()

    async with _async_session() as session:
        result = await session.execute(
            _select(RssItem, RssFeed.title.label("feed_title"))
            .join(RssFeed, RssItem.feed_id == RssFeed.id)
            .where(RssItem.id == item_id, RssFeed.user_id == uid)
        )
        row = result.first()

    if row is None:
        return jsonify({"error": "Article not found"}), 404

    item, feed_title = row

    conv_title = (item.title or "Article discussion")[:80]
    conv = await create_conversation(uid, title=conv_title, conversation_type="chat")

    from fabledassistant.services.rss import _fetch_full_article
    source = feed_title or "News"
    content_body = (await _fetch_full_article(item.url) if item.url else None) or (item.content or "").strip()
    seeded_text = f"**{source}**\n\n**{item.title}**"
    if content_body:
        seeded_text += f"\n\n{content_body}"
    if item.url:
        seeded_text += f"\n\nSource: {item.url}"

    async with _async_session() as session:
        msg = Message(
            conversation_id=conv.id,
            role="assistant",
            content=seeded_text,
            msg_metadata={"rss_item_ids": [item_id]},
        )
        session.add(msg)
        await session.commit()

    return jsonify({"conversation_id": conv.id}), 201