#!/usr/bin/env python3
"""
orchestrator.py — OSVauco multi-agent orchestrator using ADK 2.0 collaborative workflow.

Architecture:
  root_agent (coordinator / LLM-driven delegation)
  ├── rag_agent        — retrieves answers from private knowledge base
  ├── gcp_ops_agent    — handles GCP operations questions and script generation
  ├── memory_agent     — loads/stores long-term memory via Memory Bank
  └── farewell_agent   — session closings

Pattern: Coordinator with sub_agents list.
Requires: google-adk >= 2.0.0
"""
from __future__ import annotations

import os
import logging
from typing import Optional

from google.adk.agents import Agent
from google.adk.agents.callback_context import CallbackContext
from google.adk.tools.load_memory_tool import load_memory_tool
from google.adk.tools.preload_memory_tool import preload_memory_tool
from google.genai.types import Content, Part

logger = logging.getLogger(__name__)

PROJECT_ID = os.environ.get("PROJECT_ID", "propane-will-491900-m5")
REGION = os.environ.get("REGION", "us-central1")
RAG_CORPUS_NAME = os.environ.get("RAG_CORPUS_NAME", "")
ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.5-flash")
SUBAGENT_MODEL = os.environ.get("SUBAGENT_MODEL", "gemini-2.5-flash")

# ── Safety callbacks ──────────────────────────────────────────────────────────
BLOCKED_PATTERNS = [
    "ignore previous instructions",
    "ignore all instructions",
    "drop table",
    "system prompt",
    "jailbreak",
    "disregard safety",
]

def before_model_callback(callback_context: CallbackContext, llm_request) -> Optional[Content]:
    try:
        user_text = ""
        if llm_request.contents:
            last = llm_request.contents[-1]
            if last.parts:
                user_text = last.parts[0].text.lower()
        for pattern in BLOCKED_PATTERNS:
            if pattern in user_text:
                logger.warning("Blocked pattern detected: '%s'", pattern)
                return Content(parts=[Part(text="I cannot process that request. Please rephrase.")])
    except Exception as exc:
        logger.error("before_model_callback error: %s", exc)
    return None

def before_tool_callback(tool, args: dict, tool_context) -> Optional[dict]:
    dangerous = ["DROP", "DELETE FROM", "TRUNCATE", "--", ";"]
    for val in args.values():
        if isinstance(val, str):
            for d in dangerous:
                if d.upper() in val.upper():
                    raise ValueError(f"Tool argument rejected by safety guardrail: '{val}'")
    return None

# ── Sub-agents ────────────────────────────────────────────────────────────────
rag_agent = Agent(
    model=SUBAGENT_MODEL,
    name="rag_agent",
    description="Retrieves answers from the private OSVauco knowledge base (RAG corpus).",
    instruction=(
        "You are a knowledge retrieval specialist. "
        "Use search_knowledge_base to find relevant information and return well-cited answers."
    ),
    before_model_callback=before_model_callback,
)

gcp_ops_agent = Agent(
    model=SUBAGENT_MODEL,
    name="gcp_ops_agent",
    description="Answers GCP operations questions: scripts, IAM, billing, Cloud Run, ADK deployments.",
    instruction=(
        "You are a GCP operations expert for project propane-will-491900-m5 in us-central1. "
        "Provide accurate gcloud CLI commands, IAM patterns, and ADK deployment guidance. "
        "Always include cost-safety reminders (teardown, billing budgets)."
    ),
    before_model_callback=before_model_callback,
)

memory_agent = Agent(
    model=SUBAGENT_MODEL,
    name="memory_agent",
    description="Manages long-term memory: loads past context and stores new facts for future sessions.",
    instruction=(
        "You manage the agent's long-term memory. "
        "Use load_memory_tool to retrieve past facts and preload_memory_tool to store important new facts. "
        "Only persist high-value, factual information — not transient conversation."
    ),
    tools=[load_memory_tool, preload_memory_tool],
    before_model_callback=before_model_callback,
)

farewell_agent = Agent(
    model=SUBAGENT_MODEL,
    name="farewell_agent",
    description="Handles session closings, summaries, and goodbye messages.",
    instruction="Generate a concise, friendly session summary and closing message.",
)

# ── Root orchestrator ─────────────────────────────────────────────────────────
root_agent = Agent(
    model=ORCHESTRATOR_MODEL,
    name="oavauco_orchestrator",
    description="OSVauco root orchestrator — delegates to specialist sub-agents.",
    instruction=(
        "You are the OSVauco orchestrator for project propane-will-491900-m5. "
        "Delegate to sub-agents based on the user's intent:\n"
        "- Knowledge base questions → rag_agent\n"
        "- GCP operations, scripts, IAM, billing → gcp_ops_agent\n"
        "- Memory recall or storage → memory_agent\n"
        "- Session endings → farewell_agent\n"
        "Always synthesize sub-agent responses into a clear, concise final answer."
    ),
    sub_agents=[rag_agent, gcp_ops_agent, memory_agent, farewell_agent],
    before_model_callback=before_model_callback,
    before_tool_callback=before_tool_callback,
)