WealthySmart/backend/app/main.py

import asyncio
import json
import uuid
from time import perf_counter
from http.cookies import SimpleCookie
from contextlib import asynccontextmanager

from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from jose import JWTError, jwt
from pydantic import BaseModel

from app.agent.agent import build_agent
from app.agent.chat_audit import (
    create_chat_run,
    finish_pending_chat_run,
    mark_chat_run_failed,
    purge_expired_chat_runs_periodically,
    reset_current_chat_run,
    set_current_chat_run,
)
from app.agent.snapshot_store import PostgresSnapshotStore
from app.agent.tools import reset_session, set_session
from app.api.v1.router import api_router
from app.auth import (
    ALGORITHM,
    check_login_rate_limit,
    create_access_token,
    verify_admin_credentials,
)
from app.config import settings
from app.db import get_session, run_alembic_upgrade
from app.seed import seed_db
from app.timeutil import reset_client_timezone, set_client_timezone
from app.services.exchange_rate import refresh_rates_periodically
from app.logging import configure_structured_logging, log_event


AGENT_PATH = "/api/v1/agent/agui"
MAX_USER_AGENT_LENGTH = 512


def _request_client_details(request: Request) -> tuple[str, str | None, str | None]:
    """Return the browser client details propagated through the frontend BFF.

    Production nginx-proxy supplies X-Real-IP and the frontend forwards it to
    FastAPI. For direct local development, fall back to X-Forwarded-For and
    then the socket peer. The BFF is the only production path to the backend,
    which is not publicly exposed.
    """
    client_ip = request.headers.get("x-real-ip", "").strip()
    if not client_ip:
        forwarded = request.headers.get("x-forwarded-for", "")
        client_ip = forwarded.split(",", 1)[0].strip()
    if not client_ip:
        client_ip = request.client.host if request.client else "unknown"

    user_agent = request.headers.get("user-agent", "").strip()[:MAX_USER_AGENT_LENGTH] or None
    if not user_agent:
        return client_ip, None, None
    ua = user_agent.lower()
    if "edg/" in ua or "edgios/" in ua:
        browser = "Edge"
    elif "opr/" in ua or "opera" in ua:
        browser = "Opera"
    elif "firefox/" in ua or "fxios/" in ua:
        browser = "Firefox"
    elif "chrome/" in ua or "crios/" in ua:
        browser = "Chrome"
    elif "safari/" in ua:
        browser = "Safari"
    else:
        browser = "Other"
    return client_ip, user_agent, browser


def _pair_orphan_tool_calls(messages: list) -> list:
    """Inject synthetic tool responses for any assistant tool_calls that have
    no matching tool message.  OpenAI rejects histories where a tool_calls
    entry is not immediately followed by the corresponding tool response."""
    out: list = []
    pending: list[str] = []

    def flush():
        for call_id in pending:
            out.append({"role": "tool", "tool_call_id": call_id, "content": ""})
        pending.clear()

    for msg in messages:
        role = msg.get("role", "")
        if role == "tool":
            call_id = msg.get("tool_call_id") or msg.get("toolCallId")
            if call_id and call_id in pending:
                pending.remove(call_id)
            out.append(msg)
            continue
        if role == "assistant":
            flush()
            out.append(msg)
            for tc in msg.get("tool_calls") or msg.get("toolCalls") or []:
                tc_id = tc.get("id") if isinstance(tc, dict) else None
                if tc_id:
                    pending.append(tc_id)
            continue
        flush()
        out.append(msg)

    flush()
    return out


def _drop_stale_client_history(session, thread_id, messages: list) -> list:
    """Conversation history is server-owned (the AG-UI snapshot store).

    A request that carries assistant/tool history for a thread with NO
    stored snapshot comes from a stale client — typically a second tab that
    was open when 'Nueva conversación' cleared the thread. Trusting it
    would re-persist the cleared conversation ('chunks' resurrecting after
    every clear). Keep only the last user turn; MAF starts the thread
    fresh from there.
    """
    if not thread_id or not messages:
        return messages
    has_history = any(m.get("role") in ("assistant", "tool") for m in messages)
    if not has_history:
        return messages
    from sqlmodel import select

    from app.models.models import ChatThreadSnapshot

    row = session.exec(
        select(ChatThreadSnapshot.id).where(
            ChatThreadSnapshot.thread_id == str(thread_id)
        )
    ).first()
    if row is not None:
        return messages  # thread exists; history is legitimate
    last_user = None
    for i in range(len(messages) - 1, -1, -1):
        if messages[i].get("role") == "user":
            last_user = i
            break
    return messages[last_user:] if last_user is not None else []


@asynccontextmanager
async def lifespan(app: FastAPI):
    configure_structured_logging()
    run_alembic_upgrade()
    seed_db()
    rate_refresh_task = asyncio.create_task(refresh_rates_periodically())
    chat_audit_cleanup_task = asyncio.create_task(purge_expired_chat_runs_periodically())
    try:
        yield
    finally:
        rate_refresh_task.cancel()
        chat_audit_cleanup_task.cancel()
        try:
            await rate_refresh_task
        except asyncio.CancelledError:
            pass
        try:
            await chat_audit_cleanup_task
        except asyncio.CancelledError:
            pass


app = FastAPI(title="WealthySmart API", version="0.1.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.cors_origins_list,
    allow_credentials=True,
    allow_methods=["GET", "POST", "PATCH", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["content-type", "authorization"],
)


async def log_api_request(request: Request, call_next):
    """Emit redacted, one-line diagnostics for API requests only.

    This deliberately records no request/response bodies, headers, tokens,
    passwords, or financial payloads. Full assistant content is retained only
    in the bounded chat-run audit table.
    """
    request_id = uuid.uuid4().hex
    request.state.request_id = request_id
    client_ip, user_agent, browser = _request_client_details(request)
    request.state.client_ip = client_ip
    request.state.user_agent = user_agent
    request.state.browser = browser
    started = perf_counter()
    status_code = 500
    is_logged_api = request.url.path.startswith("/api/") and request.url.path != "/api/health"

    def emit(error_type: str | None = None) -> None:
        if not is_logged_api:
            return
        route = request.scope.get("route")
        log_event(
            "api_request",
            request_id=request_id,
            method=request.method,
            path=getattr(route, "path", request.url.path),
            status_code=status_code,
            duration_ms=round((perf_counter() - started) * 1000, 1),
            error_type=error_type,
            client_ip=client_ip,
            user_agent=user_agent,
            browser=browser,
        )

    try:
        response = await call_next(request)
        status_code = response.status_code
        response.headers["X-Request-ID"] = request_id
        original_iterator = getattr(response, "body_iterator", None)
        if original_iterator is None:
            emit()
            return response

        async def logged_stream():
            error_type: str | None = None
            try:
                async for chunk in original_iterator:
                    yield chunk
            except Exception as exc:
                error_type = type(exc).__name__
                raise
            finally:
                emit(error_type)

        response.body_iterator = logged_stream()
        return response
    except Exception as exc:
        emit(type(exc).__name__)
        raise


@app.middleware("http")
async def agent_auth_and_session(request: Request, call_next):
    """For the AG-UI route, validate the JWT, repair message history, and
    bind a DB session to a ContextVar so agent tools can query without going
    through Depends."""
    if not request.url.path.startswith(AGENT_PATH):
        return await call_next(request)

    if request.method == "OPTIONS":
        return await call_next(request)

    auth_header = request.headers.get("authorization", "")
    token: str | None = None
    if auth_header.lower().startswith("bearer "):
        token = auth_header.split(" ", 1)[1].strip()
    else:
        cookies = SimpleCookie()
        cookies.load(request.headers.get("cookie", ""))
        if "ws_token" in cookies:
            token = cookies["ws_token"].value

    if not token:
        return Response(status_code=401, content="Missing auth")
    try:
        payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[ALGORITHM])
        if not payload.get("sub"):
            return Response(status_code=401, content="Invalid token")
    except JWTError:
        return Response(status_code=401, content="Invalid token")

    session_gen = get_session()
    session = next(session_gen)
    token_var = set_session(session)

    chat_run_id: int | None = None
    # Repair orphan tool_calls before the MAF agent sees the message history.
    if request.method == "POST" and "application/json" in request.headers.get("content-type", ""):
        raw = await request.body()
        try:
            body = json.loads(raw)
            if isinstance(body.get("messages"), list):
                body["messages"] = _drop_stale_client_history(
                    session,
                    body.get("threadId") or body.get("thread_id"),
                    body["messages"],
                )
                body["messages"] = _pair_orphan_tool_calls(body["messages"])
                raw = json.dumps(body).encode()
            log_event(
                "agent_request",
                request_id=request.state.request_id,
                agui_run_id=body.get("runId"),
                thread_id=body.get("threadId") or body.get("thread_id"),
                message_count=len(body.get("messages") or []),
                has_user_message=any(
                    message.get("role") == "user" for message in body.get("messages") or []
                ),
                client_ip=request.state.client_ip,
                user_agent=request.state.user_agent,
                browser=request.state.browser,
            )
            # The BFF removes its `method: agent/run` envelope before
            # forwarding to MAF. A non-empty user message is the reliable
            # backend-side distinction between an agent run and hydration.
            if any(message.get("role") == "user" for message in body.get("messages") or []):
                chat_run_id = create_chat_run(
                    request_id=request.state.request_id,
                    agui_run_id=body.get("runId"),
                    thread_id=body.get("threadId") or body.get("thread_id"),
                    messages=body.get("messages") or [],
                    client_ip=request.state.client_ip,
                    user_agent=request.state.user_agent,
                    browser=request.state.browser,
                )
        except Exception as exc:
            # Parsing/repair must not prevent the agent endpoint from serving;
            # the access log will still capture a resulting error response.
            log_event(
                "chat_audit_start_failed",
                request_id=request.state.request_id,
                error_type=type(exc).__name__,
            )
            body = None
        # Starlette caches the body; replace it so call_next sees the fixed bytes.
        request._body = raw  # type: ignore[attr-defined]
    # Browser's IANA timezone, forwarded by the BFF — lets tools resolve
    # "today" wherever the user is (falls back to Costa Rica).
    tz_token = set_client_timezone(request.headers.get("x-client-timezone"))
    audit_token = set_current_chat_run(chat_run_id) if chat_run_id is not None else None
    try:
        response = await call_next(request)
        if chat_run_id is None:
            return response

        original_iterator = response.body_iterator

        async def audit_stream():
            try:
                async for chunk in original_iterator:
                    yield chunk
            except Exception as exc:
                mark_chat_run_failed(chat_run_id, f"{type(exc).__name__}: {exc}")
                raise
            else:
                # MAF normally finalizes via PostgresSnapshotStore.save(). If
                # it emitted no snapshot, do not leave a permanently-open row.
                finish_pending_chat_run(chat_run_id)

        response.body_iterator = audit_stream()
        return response
    except Exception as exc:
        if chat_run_id is not None:
            mark_chat_run_failed(chat_run_id, f"{type(exc).__name__}: {exc}")
        raise
    finally:
        if audit_token is not None:
            reset_current_chat_run(audit_token)
        reset_client_timezone(tz_token)
        reset_session(token_var)
        try:
            next(session_gen)
        except StopIteration:
            pass


# FastAPI wraps function middlewares in reverse registration order. Register
# this after the AG-UI auth middleware so it observes every API response,
# including auth failures that return before calling the inner application.
app.middleware("http")(log_api_request)


# Register app routes
app.include_router(api_router)

# Mount the AG-UI agent endpoint. The snapshot store persists the chat
# thread across page reloads; scope is a constant because this is a
# single-user app and auth is already enforced by agent_auth_and_session
# (the scope resolver only sees the AGUIRequest, which carries no identity).
add_agent_framework_fastapi_endpoint(
    app,
    build_agent(),
    AGENT_PATH,
    snapshot_store=PostgresSnapshotStore(),
    snapshot_scope_resolver=lambda _request: "default",
)


@app.get("/")
def root():
    return {"app": "WealthySmart", "version": "0.1.0"}


# ── Cookie-based auth endpoints (used by the Vite SPA) ──────────────────────

class LoginRequest(BaseModel):
    username: str
    password: str


@app.post("/api/auth/login")
def cookie_login(body: LoginRequest, request: Request, response: Response):
    check_login_rate_limit(request)
    verify_admin_credentials(body.username, body.password)
    token = create_access_token(body.username)
    response.set_cookie(
        key="ws_token",
        value=token,
        httponly=True,
        samesite="lax",
        max_age=settings.ACCESS_TOKEN_EXPIRE_MINUTES * 60,
        secure=settings.COOKIE_SECURE,
    )
    return {"ok": True}


@app.post("/api/auth/logout", status_code=204)
def cookie_logout(response: Response):
    response.delete_cookie("ws_token")


@app.get("/api/health")
def health():
    return {"ok": True}