server.py

"""
server.py — FastAPI backend (security-hardened)

Security features applied:
  - Input validation (length, charset, regex)
  - Prompt injection detection
  - AI guardrails via hardened system prompt
  - Output sanitization before client delivery
  - Rate limiting per IP (10 req / 60s)
  - Locked-down CORS (localhost only)
  - Security response headers (CSP, X-Frame-Options, etc.)
  - No stack traces exposed to clients
  - Audit logging for every request/response
  - Tool call allowlist enforcement
  - Agent loop depth cap (prevents infinite loops)
"""

import asyncio
import json
import logging
import re
from pathlib import Path

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, field_validator
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated, Optional
import operator

from tools.weather import get_weather, fetch_weekly_forecast
from tools.activities import get_activity_recommendations, generate_category_activities
from security import (
    validate_location,
    validate_profile_field,
    check_inputs_for_injection,
    build_safe_prompt,
    sanitize_llm_output,
    rate_limiter,
    fingerprint_request,
    log_request,
    log_response,
    HARDENED_SYSTEM_PROMPT,
    WEEKLY_SUMMARY_PROMPT,
    ValidationError,
)

logger = logging.getLogger("weather_agent.server")

# ── Weekly summary fallback ───────────────────────────────────────────────────

def _build_fallback_summary(days: list, location: str) -> str:
    """Generate an upbeat weekly summary from forecast data when the LLM returns nothing."""
    if not days:
        return f"Great news for {location}! Check back soon for your full weekly weather summary."

    best_day = max(days, key=lambda d: d.get("high_f", 0) - d.get("precip_mm", 0) * 2)
    rainy = [d for d in days if d.get("precip_mm", 0) > 5]
    warmest = max(days, key=lambda d: d.get("high_f", 0))

    parts = [f"It's shaping up to be a great week in {location}!"]
    parts.append(
        f"{best_day['day_name']} looks like the standout day for getting outside "
        f"with {best_day['condition'].lower()} conditions."
    )
    if rainy:
        rain_names = " and ".join(d["day_name"] for d in rainy[:2])
        parts.append(f"Keep an umbrella handy on {rain_names}.")
    parts.append(
        f"Temperatures peak around {round(warmest['high_f'])}°F on {warmest['day_name']} — "
        "make the most of it and enjoy every day!"
    )
    return " ".join(parts)


# ── App setup ─────────────────────────────────────────────────────────────────

app = FastAPI(
    title="Weather Activity Agent",
    docs_url=None,      # Disable Swagger UI in production
    redoc_url=None,     # Disable ReDoc in production
    openapi_url=None,   # Don't expose schema publicly
)

# CORS: restrict to localhost only — not open to the internet
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:8000",
        "http://127.0.0.1:8000",
    ],
    allow_methods=["GET", "POST"],
    allow_headers=["Content-Type"],
    allow_credentials=False,
)


# ── Security headers middleware ───────────────────────────────────────────────

@app.middleware("http")
async def add_security_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["X-Content-Type-Options"]    = "nosniff"
    response.headers["X-Frame-Options"]           = "DENY"
    response.headers["X-XSS-Protection"]          = "1; mode=block"
    response.headers["Referrer-Policy"]           = "no-referrer"
    response.headers["Cache-Control"]             = "no-store"
    response.headers["Content-Security-Policy"]   = (
        "default-src 'self'; "
        "script-src 'self' 'unsafe-inline'; "
        "style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; "
        "font-src https://fonts.gstatic.com; "
        "connect-src 'self'; "
        "img-src 'self' data:; "
        "frame-ancestors 'none';"
    )
    return response


# ── Rate limiting middleware ──────────────────────────────────────────────────

@app.middleware("http")
async def enforce_rate_limit(request: Request, call_next):
    if request.url.path == "/ask":
        client_ip = request.client.host or "unknown"
        if not rate_limiter.is_allowed(client_ip):
            remaining = rate_limiter.requests_remaining(client_ip)
            return HTMLResponse(
                content=json.dumps({"error": "Rate limit exceeded. Please wait before sending another request."}),
                status_code=429,
                headers={
                    "Retry-After": "60",
                    "X-RateLimit-Limit": str(rate_limiter.max_requests),
                    "X-RateLimit-Remaining": str(remaining),
                }
            )
    return await call_next(request)


# ── Agent setup ───────────────────────────────────────────────────────────────

# Allowlist: only these tool names may ever be called by the LLM
ALLOWED_TOOLS = {"get_weather", "get_activity_recommendations"}
tools = [get_weather, get_activity_recommendations]
tool_map = {t.name: t for t in tools}

# Hard cap on tool call iterations to prevent runaway loops
MAX_TOOL_ITERATIONS = 4


class AgentState(TypedDict):
    messages: Annotated[list, operator.add]
    location: str
    weather_data: dict
    final_answer: str
    tool_iterations: int


def make_llm():
    return ChatOllama(
        model="qwen3.5:9b",
        base_url="http://localhost:11434",
        temperature=0.7,
        num_predict=600,
    ).bind_tools(tools)


def build_agent():
    llm_with_tools = make_llm()

    def call_llm(state: AgentState) -> dict:
        messages = [SystemMessage(content=HARDENED_SYSTEM_PROMPT)] + state["messages"]
        response = llm_with_tools.invoke(messages)
        return {"messages": [response]}

    def run_tools(state: AgentState) -> dict:
        last_message = state["messages"][-1]
        results = []
        for tool_call in last_message.tool_calls:
            name = tool_call["name"]
            args = tool_call["args"]

            # Tool allowlist check
            if name not in ALLOWED_TOOLS:
                logger.warning(f"BLOCKED unauthorized tool call: {name!r}")
                results.append(ToolMessage(
                    content="Error: that tool is not available.",
                    tool_call_id=tool_call["id"],
                    name=name,
                ))
                continue

            # Sanitize argument types
            sanitized_args = {
                k: v if isinstance(v, (str, int, float, bool, dict, list)) else str(v)
                for k, v in args.items()
            }

            logger.info(f"Tool call: {name}({sanitized_args})")

            try:
                result = tool_map[name].invoke(sanitized_args)
            except Exception as e:
                logger.error(f"Tool {name} raised exception: {e}")
                result = {"error": "Tool execution failed."}

            results.append(ToolMessage(
                content=str(result),
                tool_call_id=tool_call["id"],
                name=name,
            ))

        return {
            "messages": results,
            "tool_iterations": state.get("tool_iterations", 0) + 1,
        }

    def should_continue(state: AgentState) -> str:
        last = state["messages"][-1]
        if state.get("tool_iterations", 0) >= MAX_TOOL_ITERATIONS:
            logger.warning("Max tool iterations reached — forcing agent to stop.")
            return "end"
        if hasattr(last, "tool_calls") and last.tool_calls:
            return "tools"
        return "end"

    def extract_final(state: AgentState) -> dict:
        content = state["messages"][-1].content
        return {"final_answer": sanitize_llm_output(content)}

    graph = StateGraph(AgentState)
    graph.add_node("llm", call_llm)
    graph.add_node("tools", run_tools)
    graph.add_node("extract", extract_final)
    graph.set_entry_point("llm")
    graph.add_conditional_edges("llm", should_continue, {"tools": "tools", "end": "extract"})
    graph.add_edge("tools", "llm")
    graph.add_edge("extract", END)
    return graph.compile()


# ── Request model ─────────────────────────────────────────────────────────────

ACTIVITY_QUESTION = "What activities do you recommend for today's weather?"

class QueryRequest(BaseModel):
    location: str
    user_name: Optional[str] = None
    likes:     Optional[str] = None
    dislikes:  Optional[str] = None

    @field_validator("location")
    @classmethod
    def location_must_be_valid(cls, v):
        try:
            return validate_location(v)
        except ValidationError as e:
            raise ValueError(str(e))

    @field_validator("user_name", "likes", "dislikes", mode="before")
    @classmethod
    def profile_fields_must_be_safe(cls, v):
        if not v:
            return v
        try:
            return validate_profile_field(v)
        except ValidationError as e:
            raise ValueError(str(e))


# ── SSE streaming endpoint ────────────────────────────────────────────────────

@app.post("/ask")
async def ask_agent(req: QueryRequest, request: Request):
    client_ip = request.client.host or "unknown"
    request_id = fingerprint_request(client_ip, req.location, ACTIVITY_QUESTION)

    async def event_stream():
        try:
            # Injection check (location only — question is hardcoded server-side)
            try:
                check_inputs_for_injection(req.location, "")
            except ValidationError as e:
                logger.warning(f"[{request_id}] Injection check failed: {e}")
                yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
                return

            log_request(request_id, client_ip, req.location, ACTIVITY_QUESTION)

            # Step 1: weather
            yield f"data: {json.dumps({'type': 'status', 'message': f'Fetching weather for {req.location}...'})}\n\n"
            await asyncio.sleep(0.05)

            weather = get_weather.invoke({"location": req.location})

            if "error" in weather:
                logger.error(f"[{request_id}] Weather tool error: {weather['error']}")
                yield f"data: {json.dumps({'type': 'error', 'message': 'Could not fetch weather for that location. Please check and try again.'})}\n\n"
                log_response(request_id, success=False)
                return

            yield f"data: {json.dumps({'type': 'weather', 'data': weather})}\n\n"
            await asyncio.sleep(0.05)

            # Step 2: activities — three smaller LLM calls, streamed as each finishes
            likes    = req.likes    or ""
            dislikes = req.dislikes or ""

            yield f"data: {json.dumps({'type': 'status', 'message': 'Generating outdoor ideas...'})}\n\n"
            outdoor = await asyncio.to_thread(
                generate_category_activities, weather, "outdoor", 4, likes, dislikes
            )
            yield f"data: {json.dumps({'type': 'activities_outdoor', 'data': outdoor})}\n\n"

            yield f"data: {json.dumps({'type': 'status', 'message': 'Generating indoor ideas...'})}\n\n"
            indoor = await asyncio.to_thread(
                generate_category_activities, weather, "indoor", 3, likes, dislikes
            )
            yield f"data: {json.dumps({'type': 'activities_indoor', 'data': indoor})}\n\n"

            water: list = []
            if weather.get("temperature_f", 0) >= 70:
                yield f"data: {json.dumps({'type': 'status', 'message': 'Generating water ideas...'})}\n\n"
                water = await asyncio.to_thread(
                    generate_category_activities, weather, "water", 2, likes, dislikes
                )

            all_acts = outdoor + indoor + water
            all_acts.sort(key=lambda a: a["score"], reverse=True)
            top_picks = [a for a in all_acts if a["score"] >= 50][:5]
            yield f"data: {json.dumps({'type': 'activities_top', 'data': top_picks})}\n\n"
            await asyncio.sleep(0.05)

            # Step 3: weekly forecast
            yield f"data: {json.dumps({'type': 'status', 'message': 'Fetching weekly forecast...'})}\n\n"
            weekly = await asyncio.to_thread(fetch_weekly_forecast, req.location)
            yield f"data: {json.dumps({'type': 'weekly', 'data': weekly})}\n\n"
            await asyncio.sleep(0.05)

            # Step 4: LLM agent — today's activity recommendation
            yield "data: " + json.dumps({"type": "status", "message": "Generating today's recommendations..."}) + "\n\n"
            await asyncio.sleep(0.05)

            top_activities = top_picks
            profile = {
                "user_name": req.user_name or "",
                "likes":     req.likes     or "",
                "dislikes":  req.dislikes  or "",
            }
            safe_prompt = build_safe_prompt(
                req.location, ACTIVITY_QUESTION,
                top_activities=top_activities,
                profile=profile,
            )
            agent = build_agent()
            result = await asyncio.to_thread(
                agent.invoke,
                {
                    "messages": [HumanMessage(content=safe_prompt)],
                    "location": req.location,
                    "weather_data": {},
                    "final_answer": "",
                    "tool_iterations": 0,
                }
            )

            final = sanitize_llm_output(result.get("final_answer", ""))
            yield f"data: {json.dumps({'type': 'answer', 'message': final})}\n\n"

            # Step 5: LLM direct call — upbeat weekly summary
            if "days" in weekly:
                yield f"data: {json.dumps({'type': 'status', 'message': 'Generating weekly summary...'})}\n\n"
                await asyncio.sleep(0.05)

                day_lines = "\n".join(
                    f"- {d['day_name']} {d['display_date']}: {d['condition']}, "
                    f"High {d['high_f']}°F / Low {d['low_f']}°F, "
                    f"Precip {d['precip_mm']}mm, Wind {d['wind_max_mph']}mph"
                    for d in weekly["days"]
                )
                weekly_user_msg = (
                    f"Location: {req.location}\n\n"
                    f"7-Day Forecast:\n{day_lines}\n\n"
                    "Write an upbeat summary of the week's weather to help someone plan their activities."
                )

                llm_direct = ChatOllama(
                    model="qwen3.5:9b",
                    base_url="http://localhost:11434",
                    temperature=0.8,
                    num_predict=600,
                )
                weekly_response = await asyncio.to_thread(
                    llm_direct.invoke,
                    [SystemMessage(content=WEEKLY_SUMMARY_PROMPT),
                     HumanMessage(content=weekly_user_msg)]
                )
                # Strip <think>...</think> tokens (qwen3 thinking mode)
                content = re.sub(r"<think>.*?</think>", "", weekly_response.content, flags=re.DOTALL).strip()
                weekly_answer = sanitize_llm_output(content)
                if not weekly_answer or weekly_answer == "Unable to generate a response.":
                    weekly_answer = _build_fallback_summary(weekly["days"], req.location)
                yield f"data: {json.dumps({'type': 'weekly_answer', 'message': weekly_answer})}\n\n"

            yield f"data: {json.dumps({'type': 'done'})}\n\n"
            log_response(request_id, success=True, answer_len=len(final))

        except Exception as e:
            logger.error(f"[{request_id}] Unhandled exception: {e}", exc_info=True)
            yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Please try again.'})}\n\n"
            log_response(request_id, success=False)

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={"X-Accel-Buffering": "no", "Cache-Control": "no-cache"},
    )


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


@app.get("/", response_class=HTMLResponse)
async def index():
    html_path = Path(__file__).parent / "static" / "index.html"
    static_dir = (Path(__file__).parent / "static").resolve()
    resolved = html_path.resolve()
    # Path traversal guard
    if not str(resolved).startswith(str(static_dir)):
        raise HTTPException(status_code=403)
    if not resolved.exists():
        raise HTTPException(status_code=404)
    return HTMLResponse(resolved.read_text(encoding="utf-8"))


if __name__ == "__main__":
    import uvicorn
    print("\n🌤️  Weather Activity Agent (Hardened)")
    print("=" * 40)
    print("Open your browser at: http://localhost:8000")
    print("Audit log: agent_audit.log")
    print("=" * 40 + "\n")
    uvicorn.run(app, host="127.0.0.1", port=8000, log_level="warning")