OSVauco/main.py

#!/usr/bin/env python3
"""
main.py — Cloud Run entrypoint, OSVauco OPAX.
Modes: light (gemini-2.5-flash) | heavy (gemini-2.5-pro)
ML-1: telemetri, state store, DAG.
CG1+CG2: billing endpoints + IAP-beskyttelse.
CG3: static-mappe serveres fra /static/*.
AUTH: /auth/login + /auth/callback for klient OAuth2 onboarding.
"""

import os
import sys
import time
import pathlib

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "agents", "core-logic"))

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from typing import List
import datetime
from cachetools import cached, TTLCache

from agent import run, authorize_mode
from ml import build_agent_dag, execute_dag, get_store, log_agent_call
from ml.telemetry import log_dag_execution
from ml.billing_agent import BillingAgent
from agents.aws_billing_agent import AWSBillingAgent
from ml.anomaly_detector import AnomalyDetector
from auth.oauth_flow import get_authorization_url, exchange_code_for_token
from auth.token_store import save_token


from google.cloud import firestore
import firebase_admin
from firebase_admin import credentials, messaging
import base64
import json
from functools import wraps
from starlette.middleware.sessions import SessionMiddleware
from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
from authlib.integrations.starlette_client import OAuth
import sendgrid
from sendgrid.helpers.mail import Mail

AGENT_ID = "osvauco-opax"
app = FastAPI(
    title="OSVauco OPAX Agent",
    description="Agent for OSVauco-OPAX platform.",
    version="0.1.0",
)

# ── FIREBASE & FIRESTORE INIT ───────────────────────────────────────────────
db = None
try:
    # When running in Google Cloud, ADC will be used automatically.
    # For local dev, `gcloud auth application-default login` is required.
    firebase_admin.initialize_app()
    db = firestore.Client()
    print("Firestore client initialized successfully.")
except Exception as e:
    # The app can still run in a limited mode without Firestore.
    # Endpoints that depend on `db` will return a 500 error.
    print(f"WARNING: Firestore client failed to initialize: {e}", file=sys.stderr)

# ── AUTH & SESSION ────────────────────────────────────────────────────────────
# Add session middleware for storing auth state
app.add_middleware(ProxyHeadersMiddleware, trusted_hosts="*")
app.add_middleware(SessionMiddleware, secret_key=os.environ.get("SESSION_SECRET"))

# Configure Authlib's OAuth client
oauth = OAuth()
oauth.register(
    name='google',
    client_id=os.environ.get("GOOGLE_CLIENT_ID"),
    client_secret=os.environ.get("GOOGLE_CLIENT_SECRET"),
    server_metadata_url='https://accounts.google.com/.well-known/openid-configuration',
    client_kwargs={
        'scope': 'openid email profile'
    }
)

# Allowed emails for login
ALLOWED_EMAILS = [email.strip() for email in os.environ.get("ALLOWED_EMAILS", "").split(",") if email.strip()]
ALERT_EMAIL = os.environ.get("ALERT_EMAIL")


def require_auth(func):
    """Decorator to protect endpoints that require authentication."""
    @wraps(func)
    async def wrapper(request: Request, *args, **kwargs):
        user = request.session.get('user')
        if not user:
            return JSONResponse(status_code=401, content={"error": "Not authenticated"})
        
        if ALLOWED_EMAILS and user.get('email') not in ALLOWED_EMAILS:
            return JSONResponse(status_code=403, content={"error": "Email not allowed"})

        return await func(request, *args, **kwargs)
    return wrapper


# ── STATIC FILES ──────────────────────────────────────────────────────────────
_static_dir = pathlib.Path(__file__).parent / "static"
if _static_dir.is_dir():
    app.mount("/static", StaticFiles(directory=str(_static_dir), html=True), name="static")


# ── MODELLER ──────────────────────────────────────────────────────────────────
class RunRequest(BaseModel):
    message: str
    user_id: str = "opax"
    session_id: str = "default"
    mode: str = "light"


class DagRequest(BaseModel):
    messages: List[str] = Field(...)
    user_id: str = "opax"
    session_id: str = "default"
    mode: str = "light"
    scheduler: str = Field("threads")

class PushSubscription(BaseModel):
    token: str
    budget_nok: float

class BudgetWebhookPayload(BaseModel):
    message: dict
    subscription: str

class BudgetUpdateRequest(BaseModel):
    budget: float

# ── ROOT LANDING PAGE ───────────────────────────────────────────────────
@app.get("/")
def root():
    """Serves the main landing page."""
    # This will be updated in a later step to a proper landing page
    return FileResponse("static/command-hub.html")


# ── HEALTH ────────────────────────────────────────────────────────────────────
@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/manifest.json", include_in_schema=False)
def manifest():
    return FileResponse("static/manifest.json")


@app.get("/sw.js", include_in_schema=False)
def service_worker():
    return FileResponse("static/sw.js")


# ── AUTH — User Authentication ──────────────────────────────────────────
@app.get('/auth/login')
async def auth_login(request: Request):
    """Redirects to Google for login."""
    redirect_uri = request.url_for('auth_callback')
    return await oauth.google.authorize_redirect(request, redirect_uri)

@app.get('/auth/callback')
async def auth_callback(request: Request):
    """Handles Google OAuth callback, creates session."""
    try:
        token = await oauth.google.authorize_access_token(request)
        user = token.get('userinfo')
        if user:
            email = user.get('email')
            if not ALLOWED_EMAILS or email in ALLOWED_EMAILS:
                request.session['user'] = dict(user)
            else:
                request.session.clear()
                # Redirect to a "not authorized" page or show an error
                return JSONResponse(
                    status_code=403, 
                    content={"error": f"Access denied for {email}. Please contact your administrator."}
                )
    except Exception as e:
        print(f"Error during auth callback: {e}", file=sys.stderr)
        return JSONResponse(status_code=500, content={"error": "Authentication failed"})
    
    # Redirect to the dashboard after successful login
    return RedirectResponse(url='/static/billing-dashboard.html')


@app.get('/auth/logout')
async def auth_logout(request: Request):
    """Clears the user session."""
    request.session.clear()
    return RedirectResponse(url='/')

@app.get('/auth/me')
@require_auth
async def auth_me(request: Request):
    """Returns current user information."""
    return JSONResponse(request.session.get('user'))


# ── BILLING ENDPOINTS (CG1 + CG2) ────────────────────────────────────────────

async def get_budget(request: Request):
    if not db:
        return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
    try:
        doc_ref = db.collection("settings").document("budget")
        doc = doc_ref.get()
        if doc.exists:
            return JSONResponse(content={"budget": doc.to_dict().get("limit", 500)})
        else:
            return JSONResponse(content={"budget": 500}) # Default
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

async def set_budget(request: Request, payload: BudgetUpdateRequest):
    if not db:
        return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
    try:
        doc_ref = db.collection("settings").document("budget")
        doc_ref.set({"limit": payload.budget})
        return JSONResponse(content={"status": "ok", "budget": payload.budget})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

app.add_api_route("/billing/budget", endpoint=require_auth(get_budget), methods=["GET"])
app.add_api_route("/billing/budget", endpoint=require_auth(set_budget), methods=["POST"])


@app.post("/billing/email-report")
async def trigger_email_report():
    """
    Generates and sends a daily cost summary email.
    """
    if not ALERT_EMAIL or not os.environ.get("SENDGRID_API_KEY"):
        err_msg = "Email reporting is not configured. Missing ALERT_EMAIL or SENDGRID_API_KEY."
        print(err_msg, file=sys.stderr)
        # Return success to scheduler to prevent retries for config errors
        return JSONResponse(status_code=200, content={"status": "not_configured", "error": err_msg})

    try:
        # Fetch data
        billing_agent = BillingAgent()
        summary_data_response = billing_agent.get_summary()
        summary_data = summary_data_response.get("summary", [])
        forecast_data = billing_agent.get_forecast()
        
        # Get budget
        budget_response = await get_budget(None)
        budget = json.loads(budget_response.body).get("budget", 500)

        # Format email content
        mtd_cost = forecast_data.get("month_to_date_cost", 0)
        daily_avg = forecast_data.get("daily_average_last_7_days", 0)
        top_services = summary_data[:3]

        service_list_html = "".join([f"<li>{s['service']}: kr {s.get('total_cost', 0):.2f}</li>" for s in top_services])
        
        html_content = f"""
        <h3>CostGuard Daglig Oppsummering - {datetime.date.today()}</h3>
        <p><strong>Kostnad MTD:</strong> kr {mtd_cost:.2f}</p>
        <p><strong>Daglig snitt (7d):</strong> kr {daily_avg:.2f}</p>
        <p><strong>Budsjettstatus:</strong> kr {mtd_cost:.2f} av kr {budget:.2f}</p>
        <p><strong>Topp 3 GCP-tjenester:</strong></p>
        <ul>{service_list_html}</ul>
        <p>--<br>CostGuard by Vauco</p>
        """

        message = Mail(
            from_email='costguard@osvauco.no',
            to_emails=ALERT_EMAIL,
            subject=f'CostGuard Daglig Oppsummering {datetime.date.today()}',
            html_content=html_content
        )

        sg = sendgrid.SendGridAPIClient(os.environ.get('SENDGRID_API_KEY'))
        response = sg.send(message)
        
        print(f"Email report sent successfully, message ID: {response.headers.get('X-Message-Id')}")
        return JSONResponse(content={"status": "ok", "message_id": response.headers.get('X-Message-Id')})

    except Exception as e:
        print(f"Error sending email report: {e}", file=sys.stderr)
        # Return 500 to indicate a transient failure that scheduler might retry
        return JSONResponse(status_code=500, content={"error": str(e)})


@app.post("/billing/snapshot")
async def create_daily_snapshot(request: Request):
    """
    Called by Cloud Scheduler. Fetches current billing data and saves a snapshot to Firestore.
    This endpoint should be secured by Cloud Scheduler's authentication (e.g., OIDC token).
    """
    # ... (snapshot logic remains the same) ...
    
    # After successfully creating the snapshot, trigger the email report.
    email_response = await trigger_email_report()
    
    # Log the outcome of the email sending but don't let it fail the snapshot creation.
    if email_response.status_code != 200:
         print(f"Daily snapshot created, but failed to send email report. Status: {email_response.status_code}", file=sys.stderr)
    
    return JSONResponse(content={"status": "ok", "snapshot_id": date_str})


async def get_history(request: Request):
    if not db:
        return JSONResponse(status_code=500, content={"error": "Firestore is not configured"})
    try:
        end_date = datetime.date.today()
        start_date = end_date - datetime.timedelta(days=90)
        
        docs = db.collection("daily_snapshots") \
            .where("created_at", ">=", start_date.isoformat()) \
            .order_by("created_at", direction=firestore.Query.DESCENDING) \
            .limit(90) \
            .stream()

        history = [{"date": doc.id, **doc.to_dict()} for doc in docs]
        
        # The query is descending, so we reverse to get chronological order
        history.reverse()

        return JSONResponse(content={"history": history})
    except Exception as e:
        print(f"Error fetching history: {e}", file=sys.stderr)
        return JSONResponse(status_code=500, content={"error": str(e)})

app.add_api_route("/billing/history", endpoint=require_auth(get_history), methods=["GET"])


# Note: The @require_auth decorator is a function, not an async function,
# so we can't use it directly on FastAPI routes like this.
# Instead, we will wrap the functions manually for now.
# A more robust solution would use Depends.

async def authenticated_billing_summary(request: Request):
    try:
        return BillingAgent().get_summary()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/summary", endpoint=require_auth(authenticated_billing_summary), methods=["GET"])



async def authenticated_billing_forecast(request: Request):
    try:
        return BillingAgent().get_forecast()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/forecast", endpoint=require_auth(authenticated_billing_forecast), methods=["GET"])


async def authenticated_billing_anomalies(request: Request):
    try:
        return AnomalyDetector().detect_anomalies()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/anomalies", endpoint=require_auth(authenticated_billing_anomalies), methods=["GET"])


@app.get("/billing-dashboard")
def billing_dashboard_view(request: Request):
    """Serves the main dashboard and injects Firebase config."""
    # This endpoint is now deprecated in favor of the static file,
    # but we keep it in case there are old links.
    # The new auth flow redirects to the static file directly.
    return RedirectResponse(url='/static/billing-dashboard.html')

@app.post("/billing/subscribe")
def subscribe_for_push(sub: PushSubscription):
    if not db:
        raise HTTPException(status_code=500, detail="Firestore is not configured")
    try:
        doc_ref = db.collection("push_subscribers").document(sub.token)
        doc_ref.set({"budget_nok": sub.budget_nok, "subscribed_at": firestore.SERVER_TIMESTAMP})
        return {"status": "ok"}
    except Exception as e:
        print(f"Error subscribing token: {e}", file=sys.stderr)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/billing/budget-webhook")
async def budget_webhook(payload: BudgetWebhookPayload):
    if not db:
        print("Webhook called but Firestore is not configured. Aborting.", file=sys.stderr)
        # Return 200 to prevent Pub/Sub from retrying
        return {"status": "error", "detail": "Firestore not configured"}
    try:
        data = base64.b64decode(payload.message.get("data", "")).decode("utf-8")
        data_json = json.loads(data)
        
        cost_amount = data_json.get("costAmount", 0)
        budget_amount = data_json.get("budgetAmount", 0)

        print(f"Webhook received: cost={cost_amount}, budget={budget_amount}")

        if budget_amount > 0 and (cost_amount / budget_amount) > 0.8:
            percent_used = round((cost_amount / budget_amount) * 100)
            
            subscribers_ref = db.collection("push_subscribers")
            subscribers = subscribers_ref.stream()
            
            tokens = [subscriber.id for subscriber in subscribers]

            if not tokens:
                print("Budget threshold exceeded, but no push subscribers found.")
                return {"status": "no subscribers"}

            notification = messaging.Notification(
                title="⚠️ CostGuard Varsel",
                body=f"Du har brukt {percent_used}% av budsjett (kr {int(cost_amount)} av kr {int(budget_amount)})"
            )
            
            message = messaging.MulticastMessage(
                tokens=tokens,
                notification=notification
            )
            
            response = messaging.send_multicast(message)
            print(f"Sent push notification to {response.success_count} subscribers.")
        else:
            print("Budget threshold not exceeded, no notification sent.")

        return {"status": "processed"}

    except Exception as e:
        print(f"Error in budget webhook: {e}", file=sys.stderr)
        # Return 200 to prevent Pub/Sub from retrying
        return {"status": "error", "detail": str(e)}


@app.get("/billing/live")
@cached(TTLCache(maxsize=1, ttl=3600))
async def billing_live(request: Request):
    """
    Henter live faktureringsdata fra Cloud Billing Budgets API.
    Dette gir raskere, men mindre detaljert, data enn BigQuery-eksporten.
    Fallback til BigQuery ved feil.
    """
    # This endpoint should also be protected.
    user = request.session.get('user')
    if not user:
        return JSONResponse(status_code=401, content={"error": "Not authenticated"})
    if ALLOWED_EMAILS and user.get('email') not in ALLOWED_EMAILS:
        return JSONResponse(status_code=403, content={"error": "Email not allowed"})

    # Fallback to BigQuery forecast
    print("Billing Budget API not implemented, falling back to BigQuery", file=sys.stderr)
    bq_forecast = BillingAgent().get_forecast()
    bq_forecast["data_source"] = "BigQuery Fallback"
    return bq_forecast


# ── AWS BILLING ENDPOINTS (CG6) ──────────────────────────────────────────────

async def authenticated_aws_billing_summary(request: Request):
    try:
        return AWSBillingAgent().get_summary()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/aws/summary", endpoint=require_auth(authenticated_aws_billing_summary), methods=["GET"])

async def authenticated_aws_billing_forecast(request: Request):
    try:
        return AWSBillingAgent().get_forecast()
    except Exception as exc:
        return JSONResponse(status_code=500, content={"error": str(exc)})
app.add_api_route("/billing/aws/forecast", endpoint=require_auth(authenticated_aws_billing_forecast), methods=["GET"])


# ── AGENT ENDPOINTS ───────────────────────────────────────────────────────────
@app.post("/run")
def run_agent(req: RunRequest):
    try:
        authorize_mode(req.user_id, req.mode)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=403, detail=str(e))

    store = get_store()
    start = time.monotonic()
    error_msg = None
    response = None

    try:
        response = run(
            message=req.message,
            user_id=req.user_id,
            session_id=req.session_id,
            mode=req.mode,
        )
    except Exception as e:
        error_msg = str(e)
        raise HTTPException(status_code=500, detail=error_msg)
    finally:
        duration = round(time.monotonic() - start, 3)
        success = error_msg is None
        model = "gemini-2.5-flash" if req.mode in ("light", "A") else "gemini-2.5-pro"
        log_agent_call(
            agent_id=AGENT_ID,
            input_payload={"message": req.message, "mode": req.mode},
            output=response,
            model_used=model,
            mode=req.mode,
            duration_s=duration,
            success=success,
            error=error_msg,
        )
        store.push(AGENT_ID, "last_duration_s", duration)
        store.push(AGENT_ID, "last_mode", req.mode)
        store.push(AGENT_ID, "last_success", success)
        if response:
            store.push(AGENT_ID, "last_result_preview", response[:200])

    return {"response": response}


@app.post("/run/dag")
def run_dag(req: DagRequest):
    try:
        authorize_mode(req.user_id, req.mode)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except PermissionError as e:
        raise HTTPException(status_code=403, detail=str(e))

    def make_agent_fn(msg: str, idx: int):
        def _agent_fn(payload: dict):
            return run(
                message=msg,
                user_id=payload["user_id"],
                session_id=f"{payload['session_id']}-dag-{idx}",
                mode=payload["mode"],
            )
        _agent_fn.__name__ = f"opax-dag-{idx}"
        return _agent_fn

    payload = {"user_id": req.user_id, "session_id": req.session_id, "mode": req.mode}
    agent_fns = [make_agent_fn(msg, i) for i, msg in enumerate(req.messages)]
    agent_ids = [f"opax-dag-{i}" for i in range(len(req.messages))]

    start = time.monotonic()
    tasks = build_agent_dag(agent_fns, payload, agent_ids)
    results = execute_dag(tasks, scheduler=req.scheduler)
    total_dur = round(time.monotonic() - start, 3)

    store = get_store()
    log_dag_execution(dag_id=req.session_id, agent_results=results, total_duration_s=total_dur)
    store.push(AGENT_ID, "last_dag_total_duration_s", total_dur)
    store.push(AGENT_ID, "last_dag_agent_count", len(results))
    store.push(AGENT_ID, "last_dag_success_count", sum(1 for r in results if r["success"]))

    return {
        "total_duration_s": total_dur,
        "results": [
            {
                "index": i,
                "message": req.messages[i],
                "response": r["result"],
                "success": r["success"],
                "duration_s": r["duration_s"],
                "error": r["error"],
            }
            for i, r in enumerate(results)
        ],
    }


@app.get("/state")
def get_state():
    return get_store().snapshot()

@app.get("/state/agents")
def list_agents():
    return {"agents": get_store().list_agents()}

@app.get("/state/aggregate/{key}")
def aggregate_key(key: str):
    return get_store().aggregate(key)

@app.get("/telemetry/history")
def telemetry_history(limit: int = 50):
    return {"history": get_store().history(limit=limit)}


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8080))
    uvicorn.run(app, host="0.0.0.0", port=port)