#!/usr/bin/env python3
"""
agent.py — OSVauco root agent using ADK 1.x with A / A+ mode routing.

Modes:
  A   (standard) — gemini-2.0-flash, $1/task hard stop
  A+  (audit+)   — gemini-2.5-pro orchestrator, $3/task hard stop,
                    multi-agent pipeline enabled

Authorized users for A+: opax, admin

Requires: google-adk >= 1.0.0,<2.0.0
         google-cloud-aiplatform >= 1.112.0
"""

import os
import logging
from typing import Literal

from google.adk.agents import Agent

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
LOCATION   = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1")
RAG_CORPUS = os.environ.get("RAG_CORPUS", "")

# Standard (A) models
ORCHESTRATOR_MODEL   = os.environ.get("ORCHESTRATOR_MODEL",  "gemini-2.0-flash")
SUBAGENT_MODEL       = os.environ.get("SUBAGENT_MODEL",      "gemini-2.0-flash")
REASONING_MODEL      = os.environ.get("REASONING_MODEL",     "gemini-2.0-flash")

# A+ (audit+) models
HEAVY_ORCHESTRATOR   = os.environ.get("HEAVY_ORCHESTRATOR_MODEL", "gemini-2.5-pro")
HEAVY_SUBAGENT       = os.environ.get("HEAVY_SUBAGENT_MODEL",     "gemini-2.5-pro")
HEAVY_REASONING      = os.environ.get("HEAVY_REASONING_MODEL",    "gemini-2.5-flash")

# Budget hard stops (USD per task)
BUDGET_A             = float(os.environ.get("BUDGET_A_USD_PER_TASK",       "1.0"))
BUDGET_APLUS         = float(os.environ.get("HEAVY_MODE_BUDGET_USD_PER_DAY", "3.0"))

HEAVY_MODE_ALLOWED_USERS = ["opax", "admin"]

Mode = Literal["A", "A+"]


# ---------------------------------------------------------------------------
# RAG tool
# ---------------------------------------------------------------------------

rag_tool = None
if RAG_CORPUS:
    try:
        from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval
        from vertexai.preview import rag

        rag_tool = VertexAiRagRetrieval(
            name="retrieve_knowledge",
            description="Retrieve relevant documentation and context from the OSVauco knowledge base.",
            rag_resources=[rag.RagResource(rag_corpus=RAG_CORPUS)],
            similarity_top_k=10,
            vector_distance_threshold=0.6,
        )
        logger.info(f"RAG tool initialised with corpus: {RAG_CORPUS}")
    except ImportError as e:
        logger.warning(f"VertexAiRagRetrieval not available — running without RAG: {e}")
else:
    logger.warning("RAG_CORPUS env var not set — running without RAG retrieval")


# ---------------------------------------------------------------------------
# Mode helpers
# ---------------------------------------------------------------------------

def get_models_for_mode(mode: Mode) -> dict:
    """Return model config dict for the given mode."""
    if mode == "A+":
        return {
            "orchestrator": HEAVY_ORCHESTRATOR,
            "subagent":     HEAVY_SUBAGENT,
            "reasoning":    HEAVY_REASONING,
            "budget_usd":   BUDGET_APLUS,
            "multi_agent":  True,
        }
    return {
        "orchestrator": ORCHESTRATOR_MODEL,
        "subagent":     SUBAGENT_MODEL,
        "reasoning":    REASONING_MODEL,
        "budget_usd":   BUDGET_A,
        "multi_agent":  False,
    }


def authorize_mode(user_id: str, mode: Mode) -> None:
    """Raise PermissionError if user is not authorized for A+ mode."""
    if mode == "A+" and user_id not in HEAVY_MODE_ALLOWED_USERS:
        raise PermissionError(
            f"User '{user_id}' is not authorized for A+ mode. "
            f"Authorized: {HEAVY_MODE_ALLOWED_USERS}"
        )


# ---------------------------------------------------------------------------
# Agent factory
# ---------------------------------------------------------------------------

def build_agent(mode: Mode = "A") -> Agent:
    """Build and return an ADK Agent configured for the given mode."""
    models = get_models_for_mode(mode)
    mode_label = "A+ (audit+)" if mode == "A+" else "A (standard)"

    instruction = (
        f"Du er OPAX — OSVauco AI-agent. "
        f"Modus: {mode_label}. Budsjettgrense per oppgave: ${models['budget_usd']}. "
        "Bruk retrieve_knowledge-verktøyet for å hente dokumentasjon og kontekst. "
        "Foretrekk alltid dokumenterte svar fremfor spekulasjon. "
        "Svar alltid på norsk (bokmål) med mindre brukeren eksplisitt ber om et annet språk."
    )

    return Agent(
        model=models["orchestrator"],
        name="opax_agent",
        description=f"OPAX — OSVauco enterprise agent [{mode_label}]",
        instruction=instruction,
        tools=[rag_tool] if rag_tool else [],
    )


# ---------------------------------------------------------------------------
# Default root agent (A mode) — used by ADK runner and Cloud Run
# ---------------------------------------------------------------------------

root_agent = build_agent(mode="A")


# ---------------------------------------------------------------------------
# run() — programmatic entry point for Cloud Run /run endpoint
# ---------------------------------------------------------------------------

def run(
    message: str,
    user_id: str = "opax",
    session_id: str = "default",
    mode: str = "A",
) -> str:
    """
    Handle a single request.

    Args:
        message:    User message / query.
        user_id:    Caller identity. A+ requires 'opax' or 'admin'.
        session_id: Session identifier for conversation continuity.
        mode:       'A' (standard) or 'A+' (audit+).

    Returns:
        Agent response as a string.

    Raises:
        PermissionError: If user_id is not authorized for the requested mode.
        ValueError:      If mode is not 'A' or 'A+'.
    """
    if mode not in ("A", "A+"):
        raise ValueError(f"Invalid mode '{mode}'. Must be 'A' or 'A+'.")

    authorize_mode(user_id, mode)

    agent = build_agent(mode=mode)
    response = agent.run(
        message=message,
        session_id=session_id,
        user_id=user_id,
    )
    return response.text if hasattr(response, "text") else str(response)


# ---------------------------------------------------------------------------
# CLI smoke test
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys

    query          = sys.argv[1] if len(sys.argv) > 1 else "Hva er OPAX A+ mode?"
    requested_mode = sys.argv[2] if len(sys.argv) > 2 else "A"

    print(f"Mode : {requested_mode}")
    print(f"Query: {query}")
    print("-" * 60)
    print(run(message=query, user_id="opax", mode=requested_mode))