from __future__ import annotations
from pydantic import BaseModel, Field
from typing import List, Literal, Union

from google.adk.agents import Agent
from google.adk.tools.agent_tool import AgentTool
from pydantic import BaseModel, Field


class Cat(BaseModel):
    """Instruction for creating a new survey."""
    type: Literal["cat"] = Field(default="cat", description="The type of animal")
    colour_of_fur: str = Field(
        default=None,
        description="The color of the cat's fur",
    )


class Dog(BaseModel):
    """Instruction for opening an existing survey."""
    type: Literal["dog"] = Field(default="dog", description="The type of animal")
    wagginess_of_tail: str = Field(
        default=None,
        description="The wagging behavior of the dog's tail",
    )

Animal = Union[Dog, Cat]


class ResponseMessage(BaseModel):
    summary: str = Field(description="A brief summary of the animals.")
    do: List[Animal] = Field(description="A list of animals to be created.")

class DogList(BaseModel):
    dogs: List[Dog]  = Field(description="List of dogs")

class CatList(BaseModel):
    cats: List[Cat] = Field(description="List of cats")

model='gemini-2.5-flash'

dog_agent = Agent(
    name="dog_agent",
    model=model,
    description="Gets information about dogs.",
    instruction=f"""
    You're job is to provide some information about breeds of dogs.

    Dog 1 is a golden retriever
    Dog 2 is a pug
    Dog 3 is a beagle
    Dog 4 is a poodle
    Dog 5 is a german shepherd
    Dog 6 is a labrador
    """,
    output_schema=DogList,
)

cat_agent = Agent(
    name="cat_agent",
    model=model,
    description="Gets information about cats.",
    instruction=f"""
    You're job is to provide some information about pedigrees of cats.

    Cat 1 is a siamese with brown fur
    Cat 2 is a persian with white fur
    Cat 3 is a maine coon with black fur
    Cat 4 is a ragdoll with cream fur
    Cat 5 is a sphynx with pink fur
    Cat 6 is a bengal with spotted fur
    """,
    
    output_schema=CatList,
)

root_agent = Agent(
    name="root_agent",
    model=model,
    description="Frontline agent that greets user, identifies intent, and dispatches to appropriate handlers.",
    instruction=f"""
    You're job is to get information about animals, specifically dogs and cats.

    Use the appropriate tool based on the user's request.
    Return all of the cats and dogs you get from the tools in your response.
    """,
    tools=[
        AgentTool(agent=dog_agent), AgentTool(agent=cat_agent)
        ],
    output_schema=ResponseMessage,
)