# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

%pip install --upgrade --quiet google-cloud-aiplatform[agent_engines,adk] a2a-sdk a2a-sdk[http-server] pydantic pydantic-ai fastmcp numpy python-dotenv nest-asyncio litellm

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

import os

import nest_asyncio
import vertexai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
LOCATION = "us-central1"  # @param {type: "string"}

# Create the bucket
BUCKET_NAME = f"{PROJECT_ID}-agent"
BUCKET_URI = f"gs://{BUCKET_NAME}"

# Set environment variables for ADK
os.environ["GOOGLE_CLOUD_PROJECT"] = PROJECT_ID
os.environ["GOOGLE_CLOUD_LOCATION"] = LOCATION
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "TRUE"

# For notebook async support
nest_asyncio.apply()

# Initiate the client
client = vertexai.Client(project=PROJECT_ID, location=LOCATION)

import asyncio

# General
import os
import random
import threading
import time
import warnings
from datetime import datetime, timedelta
from pathlib import Path

import httpx
import numpy as np
import uvicorn

warnings.filterwarnings("ignore")

# Agent deployment
import vertexai
from a2a.client.client import ClientConfig as A2AClientConfig
from a2a.client.client_factory import ClientFactory as A2AClientFactory
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.apps import A2AStarletteApplication
from a2a.server.events import EventQueue
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
from a2a.types import (
    AgentSkill,
    TaskState,
    TextPart,
    TransportProtocol,
    UnsupportedOperationError,
)
from a2a.types import TransportProtocol as A2ATransport
from a2a.utils import new_agent_text_message
from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH
from a2a.utils.errors import ServerError
from google.adk import Runner
from google.adk.a2a.executor.a2a_agent_executor import (
    A2aAgentExecutor,
    A2aAgentExecutorConfig,
)
from google.adk.agents import LlmAgent
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent

# Adk agent
from google.adk.models.lite_llm import LiteLlm, litellm
from google.adk.sessions import InMemorySessionService
from google.adk.tools.agent_tool import AgentTool
from google.auth import default
from google.auth.credentials import Credentials
from google.auth.transport.requests import Request as AuthRequest
from google.genai import types

# Pydantic agent
from mcp.server.fastmcp import FastMCP
from pydantic_ai import Agent
from pydantic_ai.models.google import GoogleModel
from pydantic_ai.providers.google import GoogleProvider
from vertexai.preview.reasoning_engines import A2aAgent
from vertexai.preview.reasoning_engines.templates.a2a import create_agent_card

class MarketDataGenerator:
    """Generate realistic synthetic market data."""

    def __init__(self, seed: int = 42):
        random.seed(seed)
        np.random.seed(seed)

        # Base prices for common symbols
        self.base_prices = {
            "NVDA": 850.0,
            "AAPL": 185.0,
            "GOOGL": 155.0,
            "MSFT": 420.0,
            "TSLA": 245.0,
        }

    def generate_price_series(self, symbol: str, days: int = 30) -> list[dict]:
        """Generate realistic OHLCV price series."""
        base_price = self.base_prices.get(symbol, 100.0)

        prices = [base_price]
        for _ in range(days - 1):
            drift = random.uniform(-0.005, 0.01)
            shock = random.gauss(0, 0.02)
            new_price = prices[-1] * (1 + drift + shock)
            prices.append(max(new_price, 1.0))

        # Generate OHLCV data
        ohlcv_data = []
        start_date = datetime.now() - timedelta(days=days)

        for i, close in enumerate(prices):
            date = start_date + timedelta(days=i)
            intraday_range = close * random.uniform(0.01, 0.03)
            open_price = close + random.uniform(-intraday_range / 2, intraday_range / 2)
            high = max(open_price, close) + random.uniform(0, intraday_range)
            low = min(open_price, close) - random.uniform(0, intraday_range)
            volume = int(random.uniform(50_000_000, 150_000_000))

            ohlcv_data.append(
                {
                    "date": date.strftime("%Y-%m-%d"),
                    "open": round(open_price, 2),
                    "high": round(high, 2),
                    "low": round(low, 2),
                    "close": round(close, 2),
                    "volume": volume,
                }
            )

        return ohlcv_data

    def _calculate_rsi(self, prices: list[float], period: int = 14) -> float:
        """Calculate RSI indicator."""
        if len(prices) < period + 1:
            return 50.0

        deltas = np.diff(prices[-period - 1 :])
        gains = deltas.copy()
        losses = deltas.copy()
        gains[gains < 0] = 0
        losses[losses > 0] = 0
        losses = abs(losses)

        avg_gain = np.mean(gains) if len(gains) > 0 else 0
        avg_loss = np.mean(losses) if len(losses) > 0 else 0.01

        rs = avg_gain / avg_loss if avg_loss != 0 else 100
        rsi = 100 - (100 / (1 + rs))

        return rsi

    def _calculate_macd(self, prices: list[float]) -> tuple:
        """Calculate MACD and signal line."""
        if len(prices) < 26:
            return (0.0, 0.0)

        # Simplified MACD calculation
        fast_ema = np.mean(prices[-12:])
        slow_ema = np.mean(prices[-26:])
        macd = fast_ema - slow_ema
        signal = macd * 0.9

        return (macd, signal)

market_generator = MarketDataGenerator()

# Initialize MCP server for Bear Agent tools
bear_mcp = FastMCP("bear-agent-tools")


@bear_mcp.tool()
async def risk_scanner(symbol: str) -> str:
    """Scan for potential downside risks and warning signals.

    Args:
        symbol: Stock symbol to analyze (e.g., NVDA, AAPL)

    Returns:
        Risk analysis report
    """
    # Generate market data and calculate technical indicators
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    # Calculate overall risk score
    risk_score = np.random.uniform(40, 75)

    # Identify specific risks based on technical indicators
    risks = []

    if rsi > 70:  # Overbought condition
        risks.append(
            {
                "risk": "Overbought Conditions",
                "severity": "HIGH",
                "description": f"RSI at {rsi:.1f} indicates potential pullback",
                "impact": "-5% to -10%",
            }
        )

    if len(risks) == 0:  # Default risk if no technical signals
        risks.append(
            {
                "risk": "Valuation Concerns",
                "severity": "MEDIUM",
                "description": "P/E ratio elevated vs historical average",
                "impact": "-10% to -15%",
            }
        )

    # Format comprehensive risk report
    result = f"""
RISK ANALYSIS FOR {symbol}
{"=" * 40}
Current Price: ${current_price}
Risk Score: {risk_score:.1f}/100
Risk Level: {"HIGH" if risk_score > 60 else "MEDIUM"}

Identified Risks:
"""

    for risk in risks:
        result += f"\n[{risk['severity']}] {risk['risk']}"
        result += f"\n   {risk['description']}"
        result += f"\n   Potential Impact: {risk['impact']}\n"

    return result


@bear_mcp.tool()
async def divergence_detector(symbol: str) -> str:
    """Detect bearish divergences and technical weakness.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Divergence analysis report
    """
    # Generate prices
    prices = market_generator.generate_price_series(symbol, days=30)
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    # Calculate divergence score
    divergence_score = np.random.uniform(30, 70)

    result = f"""
DIVERGENCE ANALYSIS FOR {symbol}
{"=" * 40}
Divergence Score: {divergence_score:.1f}/100
RSI: {rsi:.1f}

Detected Divergences:
• RSI Bearish Divergence
   Price making highs but RSI not confirming
   Confidence: 75%

• Volume Divergence
   Declining volume on advances
   Confidence: 70%
"""

    return result


@bear_mcp.tool()
async def exit_signal_monitor(symbol: str) -> str:
    """Monitor for distribution patterns and exit signals.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Exit signal analysis
    """
    # Generate prices
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    # Calculate stop losses
    stop_aggressive = round(current_price * 0.95, 2)
    stop_moderate = round(current_price * 0.93, 2)

    result = f"""
EXIT SIGNAL MONITOR FOR {symbol}
{"=" * 40}
Current Price: ${current_price}

Exit Signals:
[MED] Distribution Pattern
   Heavy selling on up days
   Action: Reduce position size

Stop Loss Recommendations:
   Aggressive: ${stop_aggressive} (-5%)
   Moderate: ${stop_moderate} (-7%)
"""

    return result

# Define the Bear Agent's personality and role
bear_system_prompt = (
    "You are a cautious risk analyst focused on identifying potential downside catalysts, "
    "warning signals, and protective strategies. You prioritize capital preservation. "
    "Use the available MCP tools to analyze market risks comprehensively."
)

# Configure Gemini model for Vertex AI
provider = GoogleProvider(vertexai=True)
model = GoogleModel("gemini-2.5-flash", provider=provider)

# Create Pydantic AI agent with MCP tools
bear_agent = Agent(
    model=model,
    system_prompt=bear_system_prompt,
    tools=[risk_scanner, divergence_detector, exit_signal_monitor],
    retries=2,
)

async def test_bear_agent():
    # Test query for risk analysis
    query = "Analyze the risks for NVDA stock"
    print(f"Query: {query}")
    print("-" * 60)

    # Run agent and get response
    result = await bear_agent.run(query)
    print("Agent Response:\n")
    print(result.output)


# Execute the test
await test_bear_agent()

# Initialize MCP server for Bull agent
bull_mcp = FastMCP("bull-agent-tools")


@bull_mcp.tool()
async def find_breakout_patterns(symbol: str) -> str:
    """Identify bullish breakout patterns and technical setups.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Breakout analysis report
    """
    # Generate prices
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    # Calculate the score
    breakout_score = np.random.uniform(55, 85)

    result = f"""
BREAKOUT PATTERN ANALYSIS FOR {symbol}
{"=" * 40}
Current Price: ${current_price}
Breakout Score: {breakout_score:.1f}/100
Momentum: {"STRONG" if breakout_score > 70 else "MODERATE"}

Bullish Patterns:
[HIGH] Resistance Breakout
   Price breaking above key resistance
   Target: ${round(current_price * 1.08, 2)} (+8%)

[MED] Ascending Triangle
   Higher lows with resistance test
   Target: ${round(current_price * 1.10, 2)} (+10%)
"""

    return result


@bull_mcp.tool()
async def momentum_screener(symbol: str) -> str:
    """Screen for stocks with strong upward momentum.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Momentum analysis report
    """
    # Generate prices
    prices = market_generator.generate_price_series(symbol, days=30)
    closes = [p["close"] for p in prices]

    # Calculate kpis
    rsi = market_generator._calculate_rsi(closes)
    momentum_score = np.random.uniform(60, 90)

    result = f"""
MOMENTUM ANALYSIS FOR {symbol}
{"=" * 40}
Momentum Score: {momentum_score:.1f}/100
Rating: {"VERY STRONG" if momentum_score > 80 else "STRONG"}
Trend: BULLISH

Momentum Factors:
• Healthy RSI at {rsi:.1f} - room to run
• MACD bullish crossover confirmed
• Volume surge - institutions accumulating
• Uptrend pattern intact
"""

    return result


@bull_mcp.tool()
async def entry_signal_detector(symbol: str) -> str:
    """Detect optimal entry points for long positions.

    Args:
        symbol: Stock symbol to analyze

    Returns:
        Entry signal analysis
    """
    # Generate prices
    prices = market_generator.generate_price_series(symbol, days=30)

    current_price = prices[-1]["close"]
    entry_quality = np.random.uniform(60, 90)

    result = f"""
ENTRY SIGNAL ANALYSIS FOR {symbol}
{"=" * 40}
Current Price: ${current_price}
Entry Quality: {entry_quality:.1f}/100

Entry Signals:
[HIGH] Pullback to Support
   Quality entry at ${round(current_price * 0.98, 2)}
   Stop Loss: ${round(current_price * 0.95, 2)}
   Risk/Reward: 1:3

Position Sizing:
   Suggested: {"75-100%" if entry_quality > 80 else "50-75%"} of planned position
"""

    return result

# Configure LiteLLM to route requests to Vertex AI
litellm.vertex_project = os.environ.get("GOOGLE_CLOUD_PROJECT")
litellm.vertex_location = "global"

# Define the Bull Agent's personality and role
bull_system_instruction = (
    "You are an optimistic market analyst focused on identifying growth opportunities, "
    "bullish patterns, and upside catalysts. You emphasize potential gains and momentum. "
    "Use the available tools to analyze market opportunities comprehensively."
)

# Create ADK agent with Claude Sonnet via LiteLLM routing
bull_agent = LlmAgent(
    name="bull_agent",
    model=LiteLlm("vertex_ai/claude-sonnet-4-5@20250929"),
    description="Optimistic analyst focused on growth opportunities and bullish signals.",
    instruction=bull_system_instruction,
    tools=[find_breakout_patterns, momentum_screener, entry_signal_detector],
)

async def test_bull_agent():
    """Test the Bull Agent locally using ADK Runner."""
    # Create ADK Runner to manage agent execution
    runner = Runner(
        app_name=bull_agent.name,
        agent=bull_agent,
        session_service=InMemorySessionService(),  # Manages conversations
    )

    # Create a session for the conversation
    session = await runner.session_service.create_session(
        app_name=bull_agent.name,
        user_id="test_user",
        session_id="test_session",
    )

    # Test query for opportunity analysis
    query = "What are the growth opportunities for AAPL stock?"
    print(f"Query: {query}")
    print("-" * 60)

    # Format message in ADK/Gemini format
    content = types.Content(role="user", parts=[types.Part(text=query)])

    # Run agent and capture final response
    final_response = None
    async for event in runner.run_async(
        session_id=session.id, user_id="test_user", new_message=content
    ):
        # Look for the final response event
        if event.is_final_response():
            final_response = event
            break

    # Extract and display the response
    if final_response and final_response.content:
        print("Agent Response:\n")
        for part in final_response.content.parts:
            if hasattr(part, "text") and part.text:
                print(part.text)


# Execute the test
await test_bull_agent()

# Create directory structure for MCP tools
mcp_tools_dir = Path("mcp_tools")
mcp_tools_dir.mkdir(exist_ok=True)

%%writefile $mcp_tools_dir/__init__.py
"""MCP Tools package for trading agents."""

from .market_data import MarketDataGenerator

__all__ = ["MarketDataGenerator"]

%%writefile $mcp_tools_dir/market_data.py
"""Market Data Generator - Creates synthetic market data for testing."""

import random
import numpy as np
from datetime import datetime, timedelta
from typing import List, Dict


class MarketDataGenerator:
    """Generate realistic synthetic market data."""

    def __init__(self, seed: int = 42):
        random.seed(seed)
        np.random.seed(seed)

        # Base prices for common symbols
        self.base_prices = {
            "NVDA": 850.0,
            "AAPL": 185.0,
            "GOOGL": 155.0,
            "MSFT": 420.0,
            "TSLA": 245.0,
        }

    def generate_price_series(self, symbol: str, days: int = 30) -> List[Dict]:
        """Generate realistic OHLCV price series."""
        base_price = self.base_prices.get(symbol, 100.0)

        prices = [base_price]
        for _ in range(days - 1):
            drift = random.uniform(-0.005, 0.01)
            shock = random.gauss(0, 0.02)
            new_price = prices[-1] * (1 + drift + shock)
            prices.append(max(new_price, 1.0))

        # Generate OHLCV data
        ohlcv_data = []
        start_date = datetime.now() - timedelta(days=days)

        for i, close in enumerate(prices):
            date = start_date + timedelta(days=i)
            intraday_range = close * random.uniform(0.01, 0.03)
            open_price = close + random.uniform(-intraday_range/2, intraday_range/2)
            high = max(open_price, close) + random.uniform(0, intraday_range)
            low = min(open_price, close) - random.uniform(0, intraday_range)
            volume = int(random.uniform(50_000_000, 150_000_000))

            ohlcv_data.append({
                "date": date.strftime("%Y-%m-%d"),
                "open": round(open_price, 2),
                "high": round(high, 2),
                "low": round(low, 2),
                "close": round(close, 2),
                "volume": volume
            })

        return ohlcv_data

    def _calculate_rsi(self, prices: List[float], period: int = 14) -> float:
        """Calculate RSI indicator."""
        if len(prices) < period + 1:
            return 50.0

        deltas = np.diff(prices[-period-1:])
        gains = deltas.copy()
        losses = deltas.copy()
        gains[gains < 0] = 0
        losses[losses > 0] = 0
        losses = abs(losses)

        avg_gain = np.mean(gains) if len(gains) > 0 else 0
        avg_loss = np.mean(losses) if len(losses) > 0 else 0.01

        rs = avg_gain / avg_loss if avg_loss != 0 else 100
        rsi = 100 - (100 / (1 + rs))

        return rsi

    def _calculate_macd(self, prices: List[float]) -> tuple:
        """Calculate MACD and signal line."""
        if len(prices) < 26:
            return (0.0, 0.0)

        # Simplified MACD calculation
        fast_ema = np.mean(prices[-12:])
        slow_ema = np.mean(prices[-26:])
        macd = fast_ema - slow_ema
        signal = macd * 0.9  # Simplified signal

        return (macd, signal)

%%writefile $mcp_tools_dir/bear_tools.py
"""Bear Agent MCP Tools - Risk analysis tools."""

import numpy as np
from mcp.server.fastmcp import FastMCP
from market_data import MarketDataGenerator

# Initialize MCP server
mcp = FastMCP("bear-agent-tools")

# Create global market data generator
market_generator = MarketDataGenerator()


@mcp.tool()
async def risk_scanner(symbol: str) -> str:
    """Scan for potential downside risks and warning signals."""
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    risk_score = np.random.uniform(40, 75)

    risks = []
    if rsi > 70:
        risks.append({
            "risk": "Overbought Conditions",
            "severity": "HIGH",
            "description": f"RSI at {rsi:.1f} indicates potential pullback",
            "impact": "-5% to -10%",
        })

    if len(risks) == 0:
        risks.append({
            "risk": "Valuation Concerns",
            "severity": "MEDIUM",
            "description": "P/E ratio elevated vs historical average",
            "impact": "-10% to -15%",
        })

    result = f"""
RISK ANALYSIS FOR {symbol}
{'='*40}
Current Price: ${current_price}
Risk Score: {risk_score:.1f}/100
Risk Level: {"HIGH" if risk_score > 60 else "MEDIUM"}

Identified Risks:
"""

    for risk in risks:
        result += f"\\n[{risk['severity']}] {risk['risk']}"
        result += f"\\n   {risk['description']}"
        result += f"\\n   Potential Impact: {risk['impact']}\\n"

    return result


@mcp.tool()
async def divergence_detector(symbol: str) -> str:
    """Detect bearish divergences and technical weakness."""
    prices = market_generator.generate_price_series(symbol, days=30)
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    divergence_score = np.random.uniform(30, 70)

    result = f"""
DIVERGENCE ANALYSIS FOR {symbol}
{'='*40}
Divergence Score: {divergence_score:.1f}/100
RSI: {rsi:.1f}

Detected Divergences:
• RSI Bearish Divergence
   Price making highs but RSI not confirming
   Confidence: 75%

• Volume Divergence
   Declining volume on advances
   Confidence: 70%
"""

    return result


@mcp.tool()
async def exit_signal_monitor(symbol: str) -> str:
    """Monitor for distribution patterns and exit signals."""
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    stop_aggressive = round(current_price * 0.95, 2)
    stop_moderate = round(current_price * 0.93, 2)

    result = f"""
EXIT SIGNAL MONITOR FOR {symbol}
{'='*40}
Current Price: ${current_price}

Exit Signals:
[MED] Distribution Pattern
   Heavy selling on up days
   Action: Reduce position size

Stop Loss Recommendations:
   Aggressive: ${stop_aggressive} (-5%)
   Moderate: ${stop_moderate} (-7%)
"""

    return result

%%writefile $mcp_tools_dir/bear_mcp_server.py
"""Bear Agent MCP Server - Risk-focused market analysis tools."""

# Import the mcp instance with all registered tools
from bear_tools import mcp

if __name__ == "__main__":
    # Run the MCP server with STDIO transport
    mcp.run(transport="stdio")

def create_bear_agent_card():
    """Create A2A Agent Card for Bear Risk Analyst."""
    # Define the agent's capabilities as A2A skills
    skills = [
        AgentSkill(
            id="risk_analysis",
            name="Risk Factor Scanner",
            description="Identifies potential downside catalysts and risk factors",
            tags=["Risk-Analysis", "Market-Analysis"],
            examples=[
                "What are the key risks for NVDA?",
                "Analyze downside catalysts for tech stocks",
            ],
        ),
        AgentSkill(
            id="divergence_detection",
            name="Divergence Detection",
            description="Finds bearish divergences and technical weakness signals",
            tags=["Technical-Analysis", "Divergence"],
            examples=[
                "Find bearish divergences in AAPL",
            ],
        ),
        AgentSkill(
            id="exit_signals",
            name="Exit Signal Monitoring",
            description="Tracks distribution patterns and exit signals",
            tags=["Exit-Strategy", "Risk-Management"],
            examples=[
                "Monitor exit signals for NVDA",
            ],
        ),
    ]

    # Create A2A agent card for capability advertisement
    return create_agent_card(
        agent_name="Bear Risk Analyst (Pydantic AI + MCP)",
        description=(
            "A cautious risk analyst powered by Pydantic AI, "
            "focused on identifying downside catalysts and warning signals."
        ),
        skills=skills,
    )


# Generate the agent card
bear_agent_card = create_bear_agent_card()

print("Bear Agent Card:")
print(f"   Name: {bear_agent_card.name}")
print(f"   Skills: {len(bear_agent_card.skills)}")

class BearAgentExecutor(AgentExecutor):
    """Agent executor for A2A integration with Bear Agent."""

    def __init__(self):
        # Agent initialized lazily to avoid pickling issues
        self.agent = None

    def _init_agent(self):
        """Initialize Pydantic AI agent with MCP tools on deployment."""
        if self.agent is None:
            import os

            import vertexai
            from pydantic_ai import Agent
            from pydantic_ai.mcp import MCPServerStdio
            from pydantic_ai.models.google import GoogleModel
            from pydantic_ai.providers.google import GoogleProvider

            # Get configuration from environment
            project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
            location = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1")

            # Initialize Vertex AI
            vertexai.init(project=project_id, location=location)

            # System prompt for Bear Agent
            bear_system_prompt = (
                "You are a cautious risk analyst focused on identifying potential downside catalysts, "
                "warning signals, and protective strategies. You prioritize capital preservation. "
                "Use the available MCP tools to analyze market risks comprehensively."
            )

            # Create provider and model
            provider = GoogleProvider(vertexai=True)
            model = GoogleModel("gemini-2.5-flash", provider=provider)

            # Configure MCP server connection
            mcp_server = MCPServerStdio(
                "python", args=["mcp_tools/bear_mcp_server.py"], timeout=60
            )

            # Create Bear Agent
            self.agent = Agent(
                model=model,
                system_prompt=bear_system_prompt,
                toolsets=[mcp_server],
                retries=3,
            )

    async def cancel(self, context: RequestContext, event_queue: EventQueue):
        # Cancellation not supported
        raise ServerError(error=UnsupportedOperationError())

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Execute Bear Agent analysis."""
        # Initialize agent if needed
        if self.agent is None:
            self._init_agent()

        # Extract user query from A2A context
        query = context.get_user_input()
        updater = TaskUpdater(event_queue, context.task_id, context.context_id)

        # Submit task if not already submitted
        if not hasattr(context, "current_task") or not context.current_task:
            await updater.submit()

        # Mark task as actively working
        await updater.start_work()

        try:
            # Update status to show progress
            await updater.update_status(
                TaskState.working, message=new_agent_text_message("Analyzing risks...")
            )

            # Run agent with user query
            result = await self.agent.run(query)

            # Extract result text from Pydantic AI response
            if hasattr(result, "output"):
                result_text = result.output
            else:
                result_text = str(result)

            # Format response
            response = f"""
BEAR RISK ANALYSIS
{"=" * 50}

{result_text}

Analysis completed
"""
            # Add result as artifact and complete task
            await updater.add_artifact([TextPart(text=response)], name="risk_analysis")
            await updater.complete()

        except Exception as e:
            # Mark task as failed on error
            await updater.update_status(
                TaskState.failed,
                message=new_agent_text_message(f"Analysis failed: {e!s}"),
            )

%%writefile $mcp_tools_dir/bull_tools.py
"""Bull Agent MCP Tools - Opportunity analysis tools."""

import numpy as np
from mcp.server.fastmcp import FastMCP
from market_data import MarketDataGenerator

# Initialize MCP server
mcp = FastMCP("bull-agent-tools")

# Create global market data generator
market_generator = MarketDataGenerator()


@mcp.tool()
async def find_breakout_patterns(symbol: str) -> str:
    """Identify bullish breakout patterns and technical setups."""
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    breakout_score = np.random.uniform(55, 85)

    result = f"""
BREAKOUT PATTERN ANALYSIS FOR {symbol}
{'='*40}
Current Price: ${current_price}
Breakout Score: {breakout_score:.1f}/100
Momentum: {"STRONG" if breakout_score > 70 else "MODERATE"}

Bullish Patterns:
[HIGH] Resistance Breakout
   Price breaking above key resistance
   Target: ${round(current_price * 1.08, 2)} (+8%)

[MED] Ascending Triangle
   Higher lows with resistance test
   Target: ${round(current_price * 1.10, 2)} (+10%)
"""

    return result


@mcp.tool()
async def momentum_screener(symbol: str) -> str:
    """Screen for stocks with strong upward momentum."""
    prices = market_generator.generate_price_series(symbol, days=30)
    closes = [p["close"] for p in prices]
    rsi = market_generator._calculate_rsi(closes)

    momentum_score = np.random.uniform(60, 90)

    result = f"""
MOMENTUM ANALYSIS FOR {symbol}
{'='*40}
Momentum Score: {momentum_score:.1f}/100
Rating: {"VERY STRONG" if momentum_score > 80 else "STRONG"}
Trend: BULLISH

Momentum Factors:
• Healthy RSI at {rsi:.1f} - room to run
• MACD bullish crossover confirmed
• Volume surge - institutions accumulating
• Uptrend pattern intact
"""

    return result


@mcp.tool()
async def entry_signal_detector(symbol: str) -> str:
    """Detect optimal entry points for long positions."""
    prices = market_generator.generate_price_series(symbol, days=30)
    current_price = prices[-1]["close"]

    entry_quality = np.random.uniform(60, 90)

    result = f"""
ENTRY SIGNAL ANALYSIS FOR {symbol}
{'='*40}
Current Price: ${current_price}
Entry Quality: {entry_quality:.1f}/100

Entry Signals:
[HIGH] Pullback to Support
   Quality entry at ${round(current_price * 0.98, 2)}
   Stop Loss: ${round(current_price * 0.95, 2)}
   Risk/Reward: 1:3

Position Sizing:
   Suggested: {"75-100%" if entry_quality > 80 else "50-75%"} of planned position
"""

    return result

%%writefile $mcp_tools_dir/bull_mcp_server.py
"""Bull Agent MCP Server - Opportunity-focused market analysis tools."""

# Import the mcp instance with all registered tools
from bull_tools import mcp

if __name__ == "__main__":
    # Run the MCP server with STDIO transport
    mcp.run(transport="stdio")

def create_bull_agent_card():
    """Create A2A Agent Card for Bull Analyst."""
    skills = [
        AgentSkill(
            id="breakout_detection",
            name="Breakout Pattern Detection",
            description="Identify bullish breakout patterns",
            tags=["technical-analysis", "breakouts"],
            examples=["Find breakout patterns for NVDA"],
        ),
        AgentSkill(
            id="momentum_screening",
            name="Momentum Screening",
            description="Screen for stocks with strong momentum",
            tags=["momentum", "screening"],
            examples=["Find high momentum tech stocks"],
        ),
        AgentSkill(
            id="entry_signals",
            name="Entry Signal Detection",
            description="Detect optimal entry points",
            tags=["entry-points", "timing"],
            examples=["When should I buy AAPL?"],
        ),
    ]

    return create_agent_card(
        agent_name="Bull Market Analyst (ADK + MCP)",
        description=(
            "An optimistic analyst powered by Google ADK, "
            "focused on growth opportunities and bullish patterns."
        ),
        skills=skills,
    )


bull_agent_card = create_bull_agent_card()

print("Bull Agent Card:")
print(f"   Name: {bull_agent_card.name}")
print(f"   Skills: {len(bull_agent_card.skills)}")

class BullAgentExecutor(AgentExecutor):
    """Agent executor for Bull Agent."""

    def __init__(self):
        self.agent = None
        self.runner = None

    def _init_agent(self):
        """Lazy initialization of the Bull Agent and ADK Runner.
        Creates the agent and runner when first needed.
        This happens on Agent Engine after deployment, not during pickling.
        """
        if self.agent is None:
            import os

            from google.adk.agents import LlmAgent
            from google.adk.models.lite_llm import LiteLlm, litellm
            from google.adk.tools.mcp_tool import StdioConnectionParams
            from google.adk.tools.mcp_tool.mcp_toolset import (
                MCPToolset,
                StdioServerParameters,
            )

            # Set project and location
            project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
            location = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")

            # Initialize Vertex AI
            vertexai.init(project=project_id, location=location)

            # Set litellm
            litellm.vertex_project = os.environ.get("GOOGLE_CLOUD_PROJECT")
            litellm.vertex_location = "global"

            # Create Bull Agent with MCP tools
            self.agent = LlmAgent(
                model=LiteLlm("vertex_ai/claude-sonnet-4-5@20250929"),
                name="bull_market_analyst",
                instruction="""You are an optimistic market analyst focused on identifying growth
                opportunities, bullish catalysts, and upside potential. Use the available MCP
                tools to analyze market opportunities comprehensively.""",
                tools=[
                    MCPToolset(
                        connection_params=StdioConnectionParams(
                            server_params=StdioServerParameters(
                                "python",
                                args=["mcp_tools/bull_mcp_server.py"],
                                timeout=60,
                            ),
                        ),
                    )
                ],
            )

        if self.runner is None:
            from google.adk import Runner
            from google.adk.sessions import InMemorySessionService

            self.runner = Runner(
                app_name=self.agent.name,
                agent=self.agent,
                session_service=InMemorySessionService(),
            )

    async def cancel(self, context: RequestContext, event_queue: EventQueue):
        raise ServerError(error=UnsupportedOperationError())

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Execute Bull Agent analysis."""
        if not context.message:
            return

        user_id = (
            context.message.metadata.get("user_id")
            if context.message and context.message.metadata
            else "a2a_user"
        )

        updater = TaskUpdater(event_queue, context.task_id, context.context_id)

        if not hasattr(context, "current_task") or not context.current_task:
            await updater.submit()

        await updater.start_work()

        query = context.get_user_input()

        try:
            await updater.update_status(
                TaskState.working,
                message=new_agent_text_message("Analyzing opportunities..."),
            )

            # Get or create session
            from google.genai import types

            session = await self.runner.session_service.get_session(
                app_name=self.runner.app_name,
                user_id=user_id,
                session_id=context.context_id,
            ) or await self.runner.session_service.create_session(
                app_name=self.runner.app_name,
                user_id=user_id,
                session_id=context.context_id,
            )

            content = types.Content(role="user", parts=[types.Part(text=query)])

            # Run ADK agent
            final_event = None
            async for event in self.runner.run_async(
                session_id=session.id, user_id=user_id, new_message=content
            ):
                if event.is_final_response():
                    final_event = event

            # Extract response
            if final_event and final_event.content and final_event.content.parts:
                response_text = "".join(
                    part.text
                    for part in final_event.content.parts
                    if hasattr(part, "text") and part.text
                )
                if response_text:
                    await updater.add_artifact(
                        [TextPart(text=response_text)],
                        name="opportunity_analysis",
                    )
                    await updater.complete()
                    return

            await updater.update_status(
                TaskState.failed,
                message=new_agent_text_message("Failed to generate response."),
                final=True,
            )

        except Exception as e:
            await updater.update_status(
                TaskState.failed,
                message=new_agent_text_message(f"Analysis failed: {e!s}"),
                final=True,
            )

# Update Bear Agent card
bear_agent_card.url = "http://localhost:8001"
bear_agent_card.preferred_transport = TransportProtocol.jsonrpc

# Update Bull Agent card
bull_agent_card.url = "http://localhost:8002"
bull_agent_card.preferred_transport = TransportProtocol.jsonrpc

def create_bull_agent_a2a_server(agent, agent_card):
    """Create an A2A server for an ADK agent.

    This wraps an ADK agent with A2A protocol handling, making it
    accessible via HTTP endpoints that follow the A2A specification.

    Args:
        agent: The ADK agent instance (LlmAgent, Agent, etc.)
        agent_card: The A2A AgentCard describing the agent's capabilities

    Returns:
        A2AStarletteApplication instance ready to serve via uvicorn
    """
    # Create ADK Runner for the agent
    # The Runner manages agent execution, sessions, and artifacts
    runner = Runner(
        app_name=agent.name,
        agent=agent,
        session_service=InMemorySessionService(),  # Manages conversation state
    )

    # Configure A2A agent executor
    # This bridges ADK agents with the A2A protocol
    config = A2aAgentExecutorConfig()
    executor = A2aAgentExecutor(runner=runner, config=config)

    # Create A2A request handler
    # Handles incoming A2A protocol requests (message:send, get_task, etc.)
    request_handler = DefaultRequestHandler(
        agent_executor=executor,
        task_store=InMemoryTaskStore(),  # Stores task state
    )

    # Create and return A2A Starlette application
    # This is the ASGI app that uvicorn will serve
    return A2AStarletteApplication(agent_card=agent_card, http_handler=request_handler)


async def run_bull_server(agent, agent_card, port):
    """Run a single agent as an A2A server on the specified port."""
    app = create_bull_agent_a2a_server(agent, agent_card)

    # Configure uvicorn server
    config = uvicorn.Config(
        app.build(),  # Build the ASGI application
        host="127.0.0.1",  # localhost
        port=port,
        log_level="warning",  # Quiet output
        loop="none",  # Use the current event loop
    )

    server = uvicorn.Server(config)
    await server.serve()

def create_bear_a2a_server(agent_card):
    """Create A2A server for Pydantic AI Bear Agent.

    Since Bear Agent uses Pydantic AI (not ADK), we create the A2A server
    directly using the BearAgentExecutor we defined earlier.
    """
    request_handler = DefaultRequestHandler(
        agent_executor=BearAgentExecutor(),
        task_store=InMemoryTaskStore(),
    )

    return A2AStarletteApplication(agent_card=agent_card, http_handler=request_handler)


async def run_bear_server(agent_card, port):
    """Run Bear Agent A2A server (Pydantic AI)."""
    app = create_bear_a2a_server(agent_card)

    config = uvicorn.Config(
        app.build(),
        host="127.0.0.1",
        port=port,
        log_level="warning",
        loop="none",
    )

    server = uvicorn.Server(config)
    await server.serve()

async def start_a2a_servers():
    """Start both Bear and Bull agents as A2A servers."""
    # Create tasks for both servers
    # Bear Agent uses Pydantic AI, so it needs custom A2A server
    # Bull Agent uses ADK, so it uses the standard ADK A2A pattern
    tasks = [
        asyncio.create_task(run_bear_server(bear_agent_card, 8001)),
        asyncio.create_task(run_bull_server(bull_agent, bull_agent_card, 8002)),
    ]

    # Give servers time to start
    await asyncio.sleep(2)

    print("   ✓ Bear Agent A2A server: http://127.0.0.1:8001 (Pydantic AI)")
    print("   ✓ Bull Agent A2A server: http://127.0.0.1:8002 (ADK)")

    # Keep servers running
    try:
        await asyncio.gather(*tasks)
    except KeyboardInterrupt:
        print("Shutting down A2A servers...")


def run_servers_in_background():
    """Run A2A servers in a background thread."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(start_a2a_servers())

# Start the A2A servers in background thread
server_thread = threading.Thread(target=run_servers_in_background, daemon=True)
server_thread.start()

# Wait for servers to be ready
time.sleep(3)

# Create remote proxy for Bear Agent
# RemoteA2aAgent discovers capabilities via the agent card endpoint
remote_bear = RemoteA2aAgent(
    name="bear_risk_analyst",
    description="Analyzes risks and warning signals",
    agent_card=f"http://localhost:8001{AGENT_CARD_WELL_KNOWN_PATH}",
)

# Create remote proxy for Bull Agent
remote_bull = RemoteA2aAgent(
    name="bull_market_analyst",
    description="Identifies growth opportunities and bullish patterns",
    agent_card=f"http://localhost:8002{AGENT_CARD_WELL_KNOWN_PATH}",
)

# Create orchestrator that coordinates both agents
trading_orchestrator = LlmAgent(
    name="trading_strategy_orchestrator",
    model="gemini-2.5-flash",
    tools=[
        AgentTool(
            agent=remote_bear,  # Wrap remote agents as tools
        ),
        AgentTool(
            agent=remote_bull,
        ),
    ],
)

# Create Runner for the orchestrator
orchestrator_runner = Runner(
    app_name=trading_orchestrator.name,
    agent=trading_orchestrator,
    session_service=InMemorySessionService(),
)

# Create session
session = await orchestrator_runner.session_service.create_session(
    app_name=trading_orchestrator.name,
    user_id="test_user",
    session_id="orchestrator_test_session",
)

# Test query
test_query = "Should I buy NVDA stock? Analyze opportunity only."

print(f"\nQuery: {test_query}")

# Run orchestrator
content = types.Content(role="user", parts=[types.Part(text=test_query)])

final_result = None
async for event in orchestrator_runner.run_async(
    session_id=session.id, user_id="test_user", new_message=content
):
    if event.is_final_response():
        if event.content and event.content.parts:
            final_result = "".join(
                part.text
                for part in event.content.parts
                if hasattr(part, "text") and part.text
            )
        break

print(f"\nFinal Result:\n{final_result}")

# Configure transport for production deployment
bear_agent_card.preferred_transport = TransportProtocol.http_json

# Wrap agent card and executor in A2A agent
bear_a2a_agent = A2aAgent(
    agent_card=bear_agent_card, agent_executor_builder=BearAgentExecutor
)

# Deploy to Vertex AI Agent Engine
deployed_bear = client.agent_engines.create(
    agent=bear_a2a_agent,
    config={
        "display_name": "Bear Risk Analyst",
        "description": bear_agent_card.description,
        "requirements": [
            "a2a-sdk",
            "google-cloud-aiplatform[agent_engines,adk]",
            "fastmcp",  # Required for MCP tools
            "pydantic",
            "pydantic-ai",  # Required for Bear Agent
            "numpy",
        ],
        "extra_packages": ["mcp_tools"],  # Include our MCP tools package
        "staging_bucket": BUCKET_URI,
    },
)

# Configure transport for production deployment
bull_agent_card.preferred_transport = TransportProtocol.http_json

# Create A2A agent
bull_a2a_agent = A2aAgent(
    agent_card=bull_agent_card, agent_executor_builder=BullAgentExecutor
)

# Deploy to Vertex AI Agent Engine
deployed_bull = client.agent_engines.create(
    agent=bull_a2a_agent,
    config={
        "display_name": "Bull Market Analyst",
        "description": bull_agent_card.description,
        "requirements": [
            "a2a-sdk",
            "google-cloud-aiplatform[agent_engines,adk]",
            "fastmcp",  # Required for MCP tools
            "numpy",
            "litellm",
        ],
        "extra_packages": ["mcp_tools"],
        "staging_bucket": BUCKET_URI,
    },
)

# Create GoogleAuth class for httpx authentication
class GoogleAuth(httpx.Auth):
    """Custom httpx Auth class for Google Cloud authentication."""

    def __init__(self) -> None:
        # Get default credentials for the current environment
        self.credentials: Credentials
        self.project: str | None
        self.credentials, self.project = default(
            scopes=["https://www.googleapis.com/auth/cloud-platform"]
        )
        self.auth_request = AuthRequest()

    def auth_flow(self, request: httpx.Request):
        """Add Authorization header to request."""
        # Refresh credentials if expired
        if not self.credentials.valid:
            self.credentials.refresh(self.auth_request)

        # Add Authorization header
        request.headers["Authorization"] = f"Bearer {self.credentials.token}"
        yield request


# Create authenticated httpx client
authenticated_client = httpx.AsyncClient(
    timeout=120,
    auth=GoogleAuth(),  # This adds authentication to ALL requests!
)

# Create client factory for A2A communication
client_config = A2AClientConfig(
    httpx_client=authenticated_client,
    streaming=False,
    polling=False,
    supported_transports=[
        A2ATransport.http_json,
    ],
)

a2a_client_factory = A2AClientFactory(config=client_config)

# Construct Vertex AI Agent Engine API endpoint
api_endpoint = f"https://{LOCATION}-aiplatform.googleapis.com"

# Get resource names from deployed agents
bear_agent_resource_name = deployed_bear.api_resource.name
bull_agent_resource_name = deployed_bull.api_resource.name

# Build A2A endpoint URLs
bear_endpoint = f"{api_endpoint}/v1beta1/{bear_agent_resource_name}/a2a"
bull_endpoint = f"{api_endpoint}/v1beta1/{bull_agent_resource_name}/a2a"

# Create remote agent proxies pointing to deployed endpoints
remote_bear = RemoteA2aAgent(
    name="bear_risk_analyst",
    description="Analyzes risks and warning signals",
    agent_card=f"{bear_endpoint}/v1/card",
    httpx_client=authenticated_client,
    a2a_client_factory=a2a_client_factory,
)

remote_bull = RemoteA2aAgent(
    name="bull_market_analyst",
    description="Identifies growth opportunities and bullish patterns",
    agent_card=f"{bull_endpoint}/v1/card",
    httpx_client=authenticated_client,
    a2a_client_factory=a2a_client_factory,
)

trading_orchestrator = LlmAgent(
    name="trading_strategy_orchestrator",
    model="gemini-2.5-flash",
    tools=[
        AgentTool(
            agent=remote_bear,
        ),
        AgentTool(
            agent=remote_bull,
        ),
    ],
)

# Create Runner for the orchestrator
orchestrator_runner = Runner(
    app_name=trading_orchestrator.name,
    agent=trading_orchestrator,
    session_service=InMemorySessionService(),
)

# Create session
session = await orchestrator_runner.session_service.create_session(
    app_name=trading_orchestrator.name,
    user_id="test_user",
    session_id="orchestrator_test_session",
)

# Test query
test_query = "Analyze the risks for NVDA stock"

print(f"\nQuery: {test_query}")

# Run orchestrator
content = types.Content(role="user", parts=[types.Part(text=test_query)])

final_result = None
async for event in orchestrator_runner.run_async(
    session_id=session.id, user_id="test_user", new_message=content
):
    if event.is_final_response():
        if event.content and event.content.parts:
            final_result = "".join(
                part.text
                for part in event.content.parts
                if hasattr(part, "text") and part.text
            )
        break

print(f"\nFinal Result:\n{final_result}")

delete_bear_agent = False
delete_bull_agent = False

if delete_bear_agent:
    client.agent_engines.delete(deployed_bear.api_resource.name, force=True)
if delete_bull_agent:
    client.agent_engines.delete(deployed_bull.api_resource.name, force=True)