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, Response
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.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",
)

# Paths exempt from IAP enforcement (health/readiness probes reach Cloud Run
# directly without the IAP-injected x-goog-authenticated-user-email header).
IAP_EXEMPT_PATHS = {"/health", "/healthz", "/readiness", "/liveness"}

@app.middleware("http")
async def require_iap(request: Request, call_next):
    if request.url.path in IAP_EXEMPT_PATHS:
        return await call_next(request)
    if not request.headers.get("x-goog-authenticated-user-email"):
        return Response(status_code=401, content="Unauthorized")
    return await call_next(request)

# ── 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'
    }
)


@app.get('/auth/login')
async def login(request: Request):
    """Redirects user to Google's OAuth 2.0 login page."""
    redirect_uri = "https://opax.vauco.no/auth/callback"
    return await oauth.google.authorize_redirect(request, redirect_uri)

@app.get('/auth/callback', name='auth')
async def auth(request: Request):
    """
    Handles the callback from Google's OAuth.
    Stores user info in session and redirects to the billing dashboard.
    """
    token = await oauth.google.authorize_access_token(request)
    user = token.get('userinfo')
    if user:
        request.session['user'] = dict(user)
    return RedirectResponse(url='/static/billing-dashboard.html')

@app.get('/auth/me')
@require_auth
async def me(request: Request):
    """Returns the authenticated user's information."""
    user = request.session.get('user')
    return JSONResponse(user)

@app.get('/auth/logout')
async def logout(request: Request):
    """Clears the user session and logs them out."""
    request.session.pop('user', None)
    return RedirectResponse(url='/static/billing-dashboard.html')


# 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."""
    return FileResponse("static/opax.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('/admin')
@require_auth
async def admin_panel(request: Request):
    return FileResponse("static/admin.html")

@app.post('/admin/create-customer')
@require_auth
async def create_customer(request: Request):
    import subprocess, shlex
    data = await request.json()
    customer_name = data.get("customer_name", "").strip()
    project_id = data.get("project_id", "").strip()
    billing_account_id = data.get("billing_account_id", "").strip()
    alert_email = data.get("alert_email", "").strip()
    container_image = data.get("container_image", "").strip()
    region = data.get("region", "europe-north1").strip()

    if not all([customer_name, project_id, billing_account_id, alert_email, container_image]):
        raise HTTPException(status_code=400, detail="Alle felt er påkrevd")

    import re
    if not re.match(r'^[a-z0-9_-]+$', customer_name):
        raise HTTPException(status_code=400, detail="customer_name kan kun inneholde a-z, 0-9, - og _")

    customer_dir = f"infrastructure/terraform/customers/{customer_name}"
    template_dir = "infrastructure/terraform/customers/_template"

    import os, shutil
    if os.path.exists(customer_dir):
        raise HTTPException(status_code=409, detail=f"Kunde {customer_name} eksisterer allerede")

    shutil.copytree(template_dir, customer_dir)

    tfvars_content = f'''customer_id           = "{customer_name}"
project_id            = "{project_id}"
region                = "{region}"
billing_account_id    = "{billing_account_id}"
alert_email           = "{alert_email}"
billing_viewer_emails = ["chris.christiansen@vauco.no", "jason.vauger@vauco.no"]
container_image       = "{container_image}"
'''
    with open(f"{customer_dir}/terraform.tfvars", "w") as f:
        f.write(tfvars_content)

    try:
        init = subprocess.run(
            ["terraform", f"-chdir={customer_dir}", "init", "-no-color"],
            capture_output=True, text=True, timeout=120
        )
        if init.returncode != 0:
            shutil.rmtree(customer_dir)
            raise HTTPException(status_code=500, detail=f"terraform init feilet: {init.stderr[-500:]}")

        apply = subprocess.run(
            ["terraform", f"-chdir={customer_dir}", "apply", "-auto-approve", "-no-color"],
            capture_output=True, text=True, timeout=600
        )
        if apply.returncode != 0:
            raise HTTPException(status_code=500, detail=f"terraform apply feilet: {apply.stderr[-500:]}")

        return {"status": "ok", "customer": customer_name, "project_id": project_id}

    except subprocess.TimeoutExpired:
        raise HTTPException(status_code=504, detail="Terraform tok for lang tid (>10 min)")


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