# Copyright 2026 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.

# @title Import Required Packages

# Standard library imports
import json
import logging
import os
import pprint
import sys
import textwrap
import time
import token
from datetime import datetime, timezone
from typing import Any, Dict, Optional
from urllib import response

# Third-party imports
import google.auth
import ipywidgets as widgets
import markdown
import requests
from google.auth.transport.requests import Request
from google.cloud import storage
from IPython.display import clear_output, display

# Initialize logger
logger = logging.getLogger(__name__)

# @title Configure Project ID and GCS Bucket of Agent
PROJECT_ID = "" #@param {type:"string"}
SKILL_GCS_BUCKET = ""  #@param {type:"string"}
ENDPOINT = "https://aiplatform.googleapis.com"
LOCATION = "global"
IS_COLAB = True

if not PROJECT_ID:
    raise ValueError("The 'PROJECT_ID'  variable is not set.")

if not SKILL_GCS_BUCKET:
    raise ValueError("The 'SKILL_GCS_BUCKET' variable is not set.")

# @title Core Code for Multi-Turn Chat with Managed Agents

if not PROJECT_ID:
    raise ValueError("The 'PROJECT_ID'  variable is not set. Repeat previous step.")

if not SKILL_GCS_BUCKET:
    raise ValueError("The 'SKILL_GCS_BUCKET' variable is not set. Repeat previous step.")

class TokenManager:
    """Manages Google Cloud OAuth 2.0 access tokens and session states."""

    def __init__(self, is_colab: bool = False):
        self.token: Optional[str] = None
        self.expiry: Optional[datetime] = None
        self.email: Optional[str] = None
        self.is_colab: bool = is_colab

    def get_token(self) -> str:
        """Retrieves the current access token, refreshing it if necessary.

        Returns:
            str: The valid OAuth 2.0 access token.
        """
        if self.token is None or self.is_expired():
            self.refresh_token()

        return self.token

    def is_expired(self) -> bool:
        """Checks if the current token has expired.

        Returns:
            bool: True if the token is expired or has no set expiry.
        """
        if self.expiry is None:
            # If we require an active token but have no expiry, force a refresh
            return True

        # Compare current UTC time against the timezone-aware expiry datetime
        return datetime.now(timezone.utc) >= self.expiry

    def _fetch_email_from_token(self) -> None:
        """Fetches the associated email for the current token via Google's Tokeninfo API."""
        if not self.token:
            return

        url = f"https://oauth2.googleapis.com/tokeninfo?access_token={self.token}"
        try:
            response = requests.get(url, timeout=10)
            if response.status_code == 200:
                data = response.json()
                # Email is only present if "userinfo.email" scope was requested
                self.email = data.get("email")
                logger.info("Token belongs to: %s", self.email)
            else:
                logger.warning("Invalid or expired token.")
        except requests.RequestException as e:
            logger.error("Failed to reach tokeninfo API: %s", e)

    def refresh_token(self) -> None:
        """Authenticates, refreshes the session, and stores a short-lived access token."""
        logger.info("****** GETTING TOKEN ******")

        if self.is_colab:
            try:
                from google.colab import auth
                auth.authenticate_user()
            except ImportError:
                logger.error("Colab environment not detected; google.colab is missing.")
                return

        # google.auth.default() handles both standard environments and Colab
        # (as long as auth.authenticate_user() was called first in Colab)
        scopes = ["https://www.googleapis.com/auth/cloud-platform"]
        credentials, _project = google.auth.default(scopes=scopes)

        # Log the credential identity for debugging
        email = getattr(credentials, "service_account_email", None)
        if email:
            logger.info("Authenticating as service account: %s", email)
        else:
            logger.info("Authenticating with credential type: %s", type(credentials).__name__)

        # Refresh the credentials to populate the token
        credentials.refresh(Request())
        self.token = credentials.token

        if credentials.expiry:
            # google-auth returns a naive datetime object in UTC; make it timezone-aware
            self.expiry = credentials.expiry.replace(tzinfo=timezone.utc)
            logger.info("Token will expire at: %s", self.expiry.strftime("%Y-%m-%d %H:%M:%S UTC"))
        else:
            self.expiry = None
            logger.info("Token does not have a set expiration time.")

        # Mask the token in stdout to prevent leaking secrets in logs
        masked_token = f"{self.token[:20]}..." if self.token else "None"
        logger.info("Obtained access token (masked): %s", masked_token)

        self._fetch_email_from_token()

class AgentFactory:
    """Factory class to manage the lifecycle of Agents via REST API."""

    def __init__(self, project_id: str, endpoint: str, location: str):
        self.project_id = project_id
        self.endpoint = endpoint
        self.location = location
        self.token_manager = TokenManager()
        self.agents: Dict[str, Dict[str, Any]] = {}

        # Populate the initial list of agents
        self.list()

    def _get_headers(self) -> Dict[str, str]:
        """Generates standard headers required for the API requests."""
        token = self.token_manager.get_token()
        return {
            "Content-Type": "application/json; charset=utf-8",
            "Authorization": f"Bearer {token}"
        }

    def _get_base_url(self) -> str:
        """Constructs the base URL for agent operations."""
        return f"{self.endpoint}/v1beta1/projects/{self.project_id}/locations/{self.location}/agents"

    def _poll_operation(self, operation: Dict[str, Any]) -> None:
        """Polls a long-running operation until it is marked as done.

        Args:
            operation (Dict[str, Any]): The initial operation payload returned by the API.
        """
        logger.info("****** POLLING OPERATION ******")

        op_name = operation.get("name")
        if not op_name:
            logger.error("Could not get operation name from response: %s", operation)
            return

        operation_status_url = f"{self.endpoint}/v1beta1/{op_name}"

        while not operation.get("done"):
            logger.info("Polling operation status...")
            time.sleep(5)  # Wait 5 seconds before checking again

            try:
                op_response = requests.get(
                    operation_status_url,
                    headers=self._get_headers(),
                    timeout=10
                )
                op_response.raise_for_status()
                operation = op_response.json()
            except requests.RequestException as e:
                logger.error("Error occurred while polling: %s", e)
                break

        logger.info("Operation complete: %s", operation)

    def create(self, agent_config: Dict[str, Any]) -> Optional[requests.Response]:
        """Creates a new agent and polls the operation until completion."""
        logger.info("****** CREATING AGENT ******")

        try:
            response = requests.post(
                self._get_base_url(),
                headers=self._get_headers(),
                json=agent_config,  # Using the 'json' kwarg automatically serializes the dict
                timeout=10
            )
            response.raise_for_status()

            logger.info("Create request submitted successfully.")

            # Poll the long-running operation
            self._poll_operation(response.json())

            # Refresh the local cache of agents
            self.list()
            return response

        except requests.RequestException as e:
            logger.error("Failed to create agent: %s", e)
            return None

    def list(self) -> Dict[str, Dict[str, Any]]:
        """Fetches a list of all agents and updates the local cache."""
        try:
            response = requests.get(
                self._get_base_url(),
                headers=self._get_headers(),
                params={"page_size": 10},
                timeout=10
            )
            response.raise_for_status()

            parsed_data = response.json()
            self.agents = {agent["id"]: agent for agent in parsed_data.get("agents", [])}
            logger.info("Successfully fetched %d agents.", len(self.agents))

        except requests.RequestException as e:
            logger.error("Failed to list agents: %s", e)

        return self.agents

    def get(self, agent_id: str) -> Optional[Dict[str, Any]]:
        """Retrieves details for a specific agent."""
        target_agent = self.agents.get(agent_id)
        if not target_agent or "name" not in target_agent:
            logger.warning("Agent ID %s not found in local cache. Call list() to refresh.", agent_id)
            return None

        try:
            response = requests.get(
                f"{self.endpoint}/v1beta1/{target_agent['name']}",
                headers=self._get_headers(),
                timeout=10
            )
            response.raise_for_status()
            return response.json()

        except requests.RequestException as e:
            logger.error("Failed to get agent %s: %s", agent_id, e)
            return None

    def delete(self, agent_id: str) -> Optional[requests.Response]:
        """Deletes a specific agent by ID."""
        target_agent = self.agents.get(agent_id)
        if not target_agent or "name" not in target_agent:
            logger.warning("Cannot delete: Agent ID %s not found in local cache.", agent_id)
            return None

        try:
            response = requests.delete(
                f"{self.endpoint}/v1beta1/{target_agent['name']}",
                headers=self._get_headers(),
                timeout=10
            )
            response.raise_for_status()

            parsed_data = response.json()
            is_done = parsed_data.get("done", False)

            if is_done:
                logger.info("Delete succeeded for agent %s.", agent_id)
                # Remove from local cache after successful deletion
                self.agents.pop(agent_id, None)
            else:
                logger.warning("Delete operation for agent %s was accepted but is not yet done.", agent_id)

            return response

        except requests.RequestException as e:
            logger.error("Failed to delete agent %s: %s", agent_id, e)
            return None

    def delete_all(self) -> None:
        """Iterates through the local cache and deletes all agents."""
        # Using list() forces a copy of keys so we don't modify the dict while iterating over it
        for agent_id in list(self.agents.keys()):
            self.delete(agent_id)

class AgentConfig:
    """Configuration builder for the World Cup demo agent."""

    def __init__(self, project_id: str, skill_gcs_bucket: str):
        self.project_id = project_id
        # Passed in explicitly rather than relying on a global SKILL_GCS_BUCKET
        self.skill_gcs_bucket = skill_gcs_bucket

    def get_config(self) -> Dict[str, Any]:
        """Returns the configuration dictionary for the agent."""

        # textwrap.dedent strips the leading indentation so your prompt string
        # remains clean and doesn't contain unnecessary whitespace.
        system_instruction = textwrap.dedent("""\
            * You are a helpful assistant focused on analysis of World Cup statistics.
            * Follow rules and guidance in the skills directory when answering questions about the 2026 World Cup.
            * Use Google Search for questions about the current World Cup in 2026.
            * Use the repository noted below for stats about the World Cup prior to 2026.

            IMPORTANT
            * When you first start up clone this repository: https://github.com/jfjelstul/worldcup
            * Save it to the local file system and use the repository to answer questions about World Cup history, teams, and players for historical questions up to 2022.
            * Use search for questions about the current World Cup in 2026.
            * Make note of repo initialization so you don't repeat this step in future questions.

            IMPORTANT
            * When you first start up also ask the user for their name, city and favorite World Cup team.
            * Save this information in a user preferences file called user_prefs.txt.
            * Always check to see if user_prefs.txt exists before asking for the user's information.

            Rule: You must always fill out the explanation parameter and narrate your actions.\
        """)

        config = {
            "id": "world-cup-agent-demo-1",
            "base_agent": "antigravity-preview-05-2026",
            "description": "A demo agent showcasing Environment and GCS use case.",
            "system_instruction": system_instruction,
            "tools": [
                {"type": "filesystem"},
                {"type": "google_search"},
            ],
            "base_environment": {
                "type": "remote",
                "sources": [
                    {
                        "type": "gcs",
                        "source": f"gs://{self.skill_gcs_bucket}",
                        "target": "./skills"
                    }
                ],
                "network": {
                    "allowlist": [
                        {"domain": "*"}
                    ]
                }
            },
        }

        return config


INTERACTION_ID = None
CALL_ID = None

class AgentHarness:
    def __init__(self, factory):
        self.factory = factory
        self.current_agent_config = None
        self.previous_interaction_id = None
        self.environment_id = None
        self.history_file = None
        self.history = {}

        # UI Widgets
        self.out_widget = None
        self.in_widget = None

        # Ensure the /logs directory exists right away
        os.makedirs("logs", exist_ok=True)

    def load_history(self):
        """Loads the saved environments and interactions from a local JSON file."""
        if self.history_file and os.path.exists(self.history_file):
            try:
                with open(self.history_file, 'r') as f:
                    return json.load(f)
            except (json.JSONDecodeError, TypeError, AttributeError):
                return {}
        return {}

    def save_state(self, prompt):
        """Saves the current environment, interaction ID, timestamps, and input snippet."""
        if self.environment_id and self.previous_interaction_id:
            now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            snippet = prompt.replace('\n', ' ')
            if len(snippet) > 30:
                snippet = snippet[:30] + "..."

            if self.environment_id not in self.history:
                self.history[self.environment_id] = {
                    "last_seen": now,
                    "interactions": {}
                }

            self.history[self.environment_id]["last_seen"] = now

            interactions_dict = self.history[self.environment_id]["interactions"]
            if self.previous_interaction_id not in interactions_dict:
                interactions_dict[self.previous_interaction_id] = {
                    "snippet": snippet,
                    "last_seen": now
                }
            else:
                interactions_dict[self.previous_interaction_id]["last_seen"] = now

            if self.history_file:
                with open(self.history_file, 'w') as f:
                    json.dump(self.history, f, indent=2)

    def handle_streaming_response(self, response, prompt):
        global INTERACTION_ID, CALL_ID

        if response.status_code != 200:
            print(f"Error: HTTP {response.status_code}")
            print(response.text)
            return

        print("\nAgent:")

        message_html = widgets.HTML(value="")
        display(message_html)

        # UI State Management
        content_blocks = {}   # Stores the final HTML/Markdown string for each index
        text_buffers = {}     # Buffers for text chunks to intelligently handle duplicates
        active_tools = {}     # Buffers for tool chunks

        def render_ui():
            # Sort by index to keep everything in strict chronological order
            full_text = "".join(content_blocks[idx] for idx in sorted(content_blocks.keys()))

            # THE FIX FOR ASCII CHARTS & TABLES IS HERE:
            html_content = markdown.markdown(full_text, extensions=['nl2br', 'fenced_code', 'tables'])
            message_html.value = f"<div style='font-family: inherit; line-height: 1.5;'>{html_content}</div>"

        for line in response.iter_lines():
            if not line:
                continue

            decoded_line = line.decode('utf-8').strip()

            if decoded_line.startswith("event:"):
                continue

            if decoded_line.startswith("data: "):
                data_str = decoded_line[6:]

                if data_str == "[DONE]":
                    print("\n--- Stream Complete ---")
                    print(f"Interaction ID: {INTERACTION_ID}")
                    self.previous_interaction_id = INTERACTION_ID
                    self.save_state(prompt)
                    break

                try:
                    event_data = json.loads(data_str)
                    event_type = event_data.get("event_type")
                    index = event_data.get("index")

                    # ---------------------------------------------------------
                    # 1. SYSTEM EVENTS
                    # ---------------------------------------------------------
                    if event_type == "interaction.start":
                        interaction_id = event_data.get("interaction", {}).get("id")
                        if interaction_id and INTERACTION_ID != interaction_id:
                            INTERACTION_ID = interaction_id

                    elif event_type == "interaction.complete":
                        interaction_data = event_data.get("interaction", {})

                        env_id = interaction_data.get("environment_id")
                        if env_id and env_id != self.environment_id:
                            self.environment_id = env_id
                            print(f"\n[Environment Attached: {self.environment_id}]")

                        usage = interaction_data.get("usage", {})
                        if usage:
                            print("\n--- Usage Statistics ---")
                            print(f"Total Tokens:  {usage.get('total_tokens')}")
                            print(f"Input Tokens:  {usage.get('total_input_tokens')} (Thought: {usage.get('total_thought_tokens', 0)})")
                            print(f"Output Tokens: {usage.get('total_output_tokens')}")

                    # ---------------------------------------------------------
                    # 2. INDEXED EVENTS (Text and Tool Calls)
                    # ---------------------------------------------------------
                    if index is not None:

                        # -- A. TOOL START --
                        if event_type == "content.start":
                            content = event_data.get("content", {})
                            if content.get("type") == "function_call":
                                call_id = content.get("id")
                                if call_id and CALL_ID != call_id:
                                    CALL_ID = call_id
                                active_tools[index] = {"name": "", "arguments": {}}

                        # -- B. STREAMING DELTAS (Tools & Text) --
                        elif event_type == "content.delta":
                            delta = event_data.get("delta", {})

                            # Stream Tool Arguments
                            if index in active_tools:
                                if "name" in delta:
                                    active_tools[index]["name"] = delta["name"]

                                if "arguments" in delta:
                                    args_chunk = delta["arguments"]
                                    if isinstance(args_chunk, dict):
                                        if not isinstance(active_tools[index]["arguments"], dict):
                                            active_tools[index]["arguments"] = {}
                                        active_tools[index]["arguments"].update(args_chunk)
                                    elif isinstance(args_chunk, str):
                                        if not isinstance(active_tools[index]["arguments"], str):
                                            active_tools[index]["arguments"] = ""
                                        active_tools[index]["arguments"] += args_chunk

                            # Stream Conversational Text
                            elif "text" in delta and index not in active_tools:
                                delta_text = delta.get("text", "")
                                if not delta_text:
                                    continue

                                if index not in text_buffers:
                                    text_buffers[index] = ""

                                # Anti-Doubling Logic:
                                if len(text_buffers[index]) > 0 and delta_text == text_buffers[index]:
                                    pass # Ignore exact catch-up duplicates
                                elif len(text_buffers[index]) > 0 and delta_text.startswith(text_buffers[index]):
                                    text_buffers[index] = delta_text # Replace with cumulative string
                                elif len(text_buffers[index]) > 0 and text_buffers[index].startswith(delta_text):
                                    pass # Ignore redundant fragments
                                else:
                                    text_buffers[index] += delta_text # Append sequential token

                                # Update the UI block and redraw
                                content_blocks[index] = text_buffers[index]
                                render_ui()

                            # Stream Citations seamlessly into the UI
                            if "annotations" in delta:
                                cit_text = "\n\n**Citations:**\n"
                                for i, annotation in enumerate(delta.get("annotations", []), start=1):
                                    title = annotation.get("title", "Unknown Source")
                                    url = annotation.get("url", "No URL provided")
                                    cit_text += f"[{i}] [{title}]({url})\n"

                                if index not in text_buffers:
                                    text_buffers[index] = ""
                                text_buffers[index] += cit_text
                                content_blocks[index] = text_buffers[index]
                                render_ui()

                        # -- C. TOOL COMPLETE (Render to UI) --
                        elif event_type == "content.stop":
                            if index in active_tools:
                                tool = active_tools.pop(index)
                                tool_name = tool.get("name", "unknown_tool")
                                args = tool.get("arguments", {})

                                if isinstance(args, str):
                                    try:
                                        args = json.loads(args) if args else {}
                                    except json.JSONDecodeError:
                                        args = {}

                                # Try to grab explanation, fallback to summary, fallback to raw payload.
                                display_text = args.get("explanation") or args.get("toolSummary") or args.get("toolAction")

                                if not display_text:
                                    clean_args = {k: v for k, v in args.items() if k not in ['explanation', 'toolSummary', 'toolAction']}
                                    display_text = f"Payload: {json.dumps(clean_args)}" if clean_args else "Executing tool..."

                                tool_ui = f"\n<br><span style='color: #6b7280;'><i>[🛠️ <b>{tool_name}</b>: {display_text}]</i></span><br>\n"

                                content_blocks[index] = tool_ui
                                render_ui()

                except json.JSONDecodeError:
                    pass

    def interact(self, prompt):
        api_url = f"{self.factory.endpoint}/v1beta1/projects/{self.factory.project_id}/locations/{self.factory.location}/interactions"

        headers = {
            "Authorization": f"Bearer {self.factory.token_manager.get_token()}",
            "Content-Type": "application/json"
        }

        agent_uri = self.current_agent_config.get("name")

        payload = {
            "agent": agent_uri,
            "stream": True,
            "background": True,
            "store": True,
            "input": [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
        }

        if self.previous_interaction_id:
            payload["previous_interaction_id"] = self.previous_interaction_id

        if self.environment_id:
            payload["environment"] = self.environment_id
        else:
            payload["environment"] = "remote"

        response = requests.post(api_url, headers=headers, json=payload, stream=True)
        self.handle_streaming_response(response, prompt)

    def multi_turn_interact(self):
        global INTERACTION_ID

        # Clear state on fresh start
        self.environment_id = None
        self.previous_interaction_id = None
        self.current_agent_config = None
        INTERACTION_ID = None

        print("--- Setup Multi-Turn Chat ---")

        # Setup Wizard (Runs in standard console before loading UI)
        agents_list = list(self.factory.agents.values())
        print("Available Agents:")
        for index, agent in enumerate(agents_list, start=1):
            description = agent.get('description', 'No description')
            print(f"{index}. {agent['id']} - {description}")

        while True:
            choice = input("\nEnter the number of the agent you want to select: ")
            try:
                selected_index = int(choice) - 1
                if 0 <= selected_index < len(agents_list):
                    self.current_agent_config = agents_list[selected_index]
                    print(f"\nSuccess! You selected: {self.current_agent_config['id']}")

                    proj = getattr(self.factory, 'project_id', 'unknown_project')
                    loc = getattr(self.factory, 'location', 'unknown_location')
                    agent_id = self.current_agent_config.get('id', 'unknown_agent')

                    file_name = f"{proj}_{loc}_{agent_id}.json"
                    self.history_file = os.path.join("logs", file_name)
                    self.history = self.load_history()
                    break
                else:
                    print(f"Invalid selection. Please enter a number between 1 and {len(agents_list)}.")
            except ValueError:
                print("Invalid input. Please enter a valid number.")

        if self.history:
            print("\nSaved Environments:")
            envs = list(self.history.keys())
            for i, env in enumerate(envs, 1):
                last_seen = self.history[env].get("last_seen", "Unknown time")
                print(f"[{i}] Environment: {env}  (Last used: {last_seen})")

            print("[0] Start completely fresh")
            env_choice = input("\nSelect Environment (or 0 to skip): ").strip()

            if env_choice.isdigit() and 0 < int(env_choice) <= len(envs):
                self.environment_id = envs[int(env_choice) - 1]
                print(f"--> Loaded Environment: {self.environment_id}")

                interaction_data = self.history[self.environment_id].get("interactions", {})
                ints = list(interaction_data.keys())

                if ints:
                    print("\nSaved Interactions:")
                    for i, interaction in enumerate(ints, 1):
                        int_snippet = interaction_data[interaction].get("snippet", "No preview")
                        print(f"[{i}] Topic: \"{int_snippet}\"")

                    print("[0] Start a NEW interaction")
                    int_choice = input("\nSelect Interaction (or 0 for new): ").strip()
                    if int_choice.isdigit() and 0 < int(int_choice) <= len(ints):
                        self.previous_interaction_id = ints[int(int_choice) - 1]
                        INTERACTION_ID = self.previous_interaction_id
                        print(f"--> Resuming Interaction: {self.previous_interaction_id}\n")
                    else:
                        print("--> Starting NEW interaction.\n")
            else:
                print("--> Starting completely fresh.\n")

        clear_output() # Clean up the setup text to make room for the UI

        # --- UI WIDGET CREATION ---
        self.out_widget = widgets.Output(layout=widgets.Layout(
            width='100%', height='400px', border='1px solid #e0e0e0', overflow_y='auto', padding='10px'
        ))

        self.in_widget = widgets.Text(
            value='',
            placeholder="Type your message and press Enter (or type 'quit' / 'reset')...",
            description='You:',
            layout=widgets.Layout(width='100%')
        )

        # Event Handler for when the user hits 'Enter'
        def on_submit(sender):
            user_input = self.in_widget.value.strip()
            self.in_widget.value = ""  # Clear the box instantly

            if not user_input:
                return

            self.in_widget.disabled = True # Prevent duplicate sends while streaming

            # The 'with' context routes all print() statements into the widget
            with self.out_widget:
                if user_input.lower() in ['quit', 'exit']:
                    print("\n[System] Conversation ended.")
                    return # Leaves the box disabled

                if user_input.lower() == 'reset':
                    print("\n" + "="*40)
                    print("[System] Session Reset. Your next message will start a fresh environment.")
                    print("="*40 + "\n")
                    self.environment_id = None
                    self.previous_interaction_id = None
                    global INTERACTION_ID
                    INTERACTION_ID = None
                    self.in_widget.disabled = False
                    return

                print(f"\n\nYou: {user_input}")
                print("-" * 40)

                # Make the API call and stream the response
                self.interact(prompt=user_input)

            self.in_widget.disabled = False # Re-enable typing

        # Bind the submission function to the Text widget
        self.in_widget.on_submit(on_submit)

        # Display the Chat UI
        display(self.out_widget, self.in_widget)
        with self.out_widget:
            print(f"Chat UI initialized for Agent: {self.current_agent_config['id']}")
            print("Type your message below and press Enter to begin.")

def check_project_settings(project_id: str) -> None:
    """Validates GCP project settings, APIs, and IAM roles for AI Platform."""

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

    # Extract Token
    printed_token = !gcloud auth application-default print-access-token
    token = printed_token[0] if printed_token else None

    if token:
        print("✅ Success on token creation.")
    else:
        print("❌ Token creation failed.")
        return

    # Extract project number
    project_info_response = !gcloud projects describe {project_id} --format="value(projectNumber)"
    if not project_info_response:
        print("❌ Failed to retrieve project number.")
        return

    project_number = project_info_response[0]
    print(f"✅ The project number for {project_id} is: {project_number}")

    # Is aiplatform enabled
    api_name = "aiplatform.googleapis.com"
    result = !gcloud services list --project={project_id} --enabled --filter="name:{api_name}" --format="value(name)"

    if any(api_name in service for service in result):
        print(f"✅ {api_name} is enabled.")
    else:
        print(f"❌ {api_name} is NOT enabled. You may need to run:")
        print(f"   !gcloud services enable {api_name} --project={project_id}")

    # Does service account have needed bindings
    sa_email = f"service-{project_number}@gcp-sa-aiplatform.iam.gserviceaccount.com"
    sa_info_response = !gcloud projects get-iam-policy {project_id} \
        --flatten="bindings[].members" \
        --filter="bindings.members:serviceAccount:{sa_email}" \
        --format="value(bindings.role)"

    if any("aiplatform.serviceAgent" in role for role in sa_info_response):
        print("✅ The service account has the aiplatform.serviceAgent role.")
    else:
        print("❌ The service account does NOT have the aiplatform.serviceAgent role.")
        print("Attempting to add role...")
        attempt_add_role = !gcloud projects add-iam-policy-binding {project_id} \
                            --member="serviceAccount:{sa_email}" \
                            --role="roles/aiplatform.serviceAgent"

    # Extract email from token
    email = None
    try:
        token_info_response = requests.get(f"https://oauth2.googleapis.com/tokeninfo?access_token={token}", timeout=10)

        if token_info_response.status_code == 200:
            data = token_info_response.json()
            # The email will only be present if the "userinfo.email" scope was requested
            email = data.get("email")
            print(f"Token belongs to: {email}")
        else:
            print("Invalid or expired token.")
            return
    except requests.RequestException as e:
        print(f"Failed to reach tokeninfo API: {e}")
        return

    if not email:
        print("Email scope not requested; cannot verify user IAM roles.")
        return

    # Does user have needed bindings
    user_roles_response = !gcloud projects get-iam-policy {project_id} \
        --flatten="bindings[].members" \
        --filter="bindings.members:user:{email}" \
        --format="value(bindings.role)"

    print(f"Raw roles found: {user_roles_response}")

    if not user_roles_response:
        print(f"❌ The user {email} has no direct project-level roles assigned.")
        return

    # Check for owner, admin, and user roles
    is_owner = any("roles/owner" in role for role in user_roles_response)
    is_ai_admin = any("aiplatform.admin" in role for role in user_roles_response)
    is_ai_user = any("aiplatform.user" in role for role in user_roles_response)

    if is_owner:
        print(f"✅ The user {email} is a project Owner, which inherits all AI Platform permissions.")
    elif is_ai_admin:
        print(f"✅ The user {email} has the explicit aiplatform.admin role.")
    elif is_ai_user:
        print(f"✅ The user {email} has the explicit aiplatform.user role (standard user access).")
    else:
        print(f"❌ The user {email} does NOT have the aiplatform.admin, aiplatform.user, or owner role.")

def test_control_plane(project_id: str, endpoint: str, location: str, delete_all: bool = False) -> None:
    factory = AgentFactory(project_id, endpoint, location)

    if delete_all:
        print("Deleting all existing agents...")
        factory.delete_all()

    agent_list = factory.list()

    print("****** AGENT LIST ******")

    if not agent_list:
        print("No agents found. Creating simple chat agent...")
        config = AgentConfig(project_id=project_id, skill_gcs_bucket=SKILL_GCS_BUCKET)
        agent_config = config.get_config()
        factory.create(agent_config)
        agent_list = factory.list()  # Refresh the list after creation

    # agent_list is a dict, so we use .items() to print both the ID and the payload
    for agent_id, agent_data in agent_list.items():
        print(f"Agent ID: {agent_id} \nDetails: {agent_data}\n")

check_project_settings(PROJECT_ID)
test_control_plane(PROJECT_ID, ENDPOINT, LOCATION, delete_all=False)

skill = r"""
# SYSTEM DIRECTIVE: WorldCup_Analyst_2026_Master

## 1. Core Identity & Persona
*   **Role:** You are the authoritative 2026 FIFA World Cup Analyst.
*   **Tone:** Objective, analytical, authoritative, and engaging. You speak like a top-tier pundit backed by a team of data scientists.
*   **Terminology:** Use standard global football terminology (pitch, match, squad, fixture, extra time). If a user says "soccer" or "roster," seamlessly understand them but maintain your professional football lexicon.
*   **Zero Hallucination Policy:** You do not hallucinate data. If a squad has not been officially announced, state clearly that it is a "projected" or "preliminary" squad. If a match result hasn't happened, do not invent a scoreline.

## 2. 2026 Tournament Framework & Baseline Context
*You must apply this framework to all answers:*
*   **Current Timeframe:** It is currently May 2026. The tournament kicks off June 11, 2026. Teams are currently locked into pre-tournament friendlies and final roster cuts.
*   **Format:** 48 teams, 104 matches, 12 groups of four.
*   **Progression:** The top two teams from each group, plus the **8 best third-place teams**, advance to the new Round of 32.
*   **Time Zones:** 16 Host cities span four North American time zones. Always convert kickoff times to the user's local time if known; otherwise, default to the local venue time and specify the time zone (e.g., PT, CT, ET).

## 3. Data Routing Protocol (Intent -> Action)
When a user asks a question, you must first classify the **Data Intent**, and then execute your search or tool call using *only* the prescribed domains for that category. Do not mix sources.

#### A. Official Logistics (Schedule, Venues, Results, Referees)
*   **Intent:** The user wants to know *when* a match is, *where* it is played, the official score, or tournament rules.
*   **Prescribed Action:** Route strictly to FIFA.
*   **Search Execution:** `site:fifa.com/tournaments/mens/worldcup/2026`
*   **Rule:** Never use third-party sites for fixture dates or kickoff times.

#### B. Advanced Statistics & Analytics (xG, Pass Completion, Heat Maps)
*   **Intent:** The user wants data on expected goals, player performance metrics, or historical tournament data.
*   **Prescribed Action:** Route strictly to FBref or Opta.
*   **Search Execution:** `site:fbref.com` OR `site:theanalyst.com`
*   **Rule:** Ignore generic sports sites for stats. Provide the data in a Markdown table.

#### C. Odds, Betting Lines & Win Probabilities
*   **Intent:** The user is asking who is favored to win, over/under goals, or outright tournament winner odds.
*   **Prescribed Action:** Route strictly to established odds aggregators.
*   **Search Execution:** `site:vegasinsider.com "world cup"` OR `site:oddschecker.com/football/world-cup`
*   **Rule:** Always state the date/time the odds were pulled, noting that "lines are subject to change." Do not provide personal betting advice.

#### D. Rosters, Injuries & Player Market Values
*   **Intent:** The user asks who made the squad, who is injured, or what a player/team is worth.
*   **Prescribed Action:** Route strictly to Transfermarkt.
*   **Search Execution:** `site:transfermarkt.us`
*   **Rule:** Distinguish clearly between "provisional" squads and "final" 26-man rosters.

#### E. News, Journalism & Tactical Breakdown
*   **Intent:** The user asks for tactical analysis, locker room news, manager quotes, or subjective match previews.
*   **Prescribed Action:** Route strictly to premium sports journalism.
*   **Search Execution:** `site:nytimes.com/athletic` OR `site:espn.com/soccer`
*   **Rule:** Synthesize the analysis objectively. Say, "According to The Athletic's tactical breakdown..."

#### F. Official Team Communications & Direct Federation Data
*   **Intent:** The user asks for official press releases, direct quotes from a national federation, or the absolute first source of a squad announcement.
*   **Prescribed Action:** Route strictly to the specific National Football Association's official domain using the exact 48-team directory below.
*   **Directory of the 48 Qualified Nations:**

| Nation | Region | Official Domain | Search Execution |
| :--- | :--- | :--- | :--- |
| **USA** | CONCACAF | ussoccer.com | `site:ussoccer.com "World Cup"` |
| **Canada** | CONCACAF | canadasoccer.com | `site:canadasoccer.com "World Cup"` |
| **Mexico** | CONCACAF | miseleccion.mx | `site:miseleccion.mx "Mundial"` |
| **Curaçao** | CONCACAF | ffk.cw | `site:ffk.cw "World Cup"` |
| **Haiti** | CONCACAF | fhfhaiti.com | `site:fhfhaiti.com "Coupe du Monde"` |
| **Panama** | CONCACAF | fepafut.com | `site:fepafut.com "Mundial"` |
| **Argentina** | CONMEBOL | afa.com.ar | `site:afa.com.ar "Mundial"` |
| **Brazil** | CONMEBOL | cbf.com.br | `site:cbf.com.br "Copa do Mundo"` |
| **Colombia** | CONMEBOL | fcf.com.co | `site:fcf.com.co "Mundial"` |
| **Ecuador** | CONMEBOL | fef.ec | `site:fef.ec "Mundial"` |
| **Paraguay** | CONMEBOL | apf.org.py | `site:apf.org.py "Mundial"` |
| **Uruguay** | CONMEBOL | auf.org.uy | `site:auf.org.uy "Mundial"` |
| **Austria** | UEFA | oefb.at | `site:oefb.at "Weltmeisterschaft"` |
| **Belgium** | UEFA | rbfa.be | `site:rbfa.be "World Cup"` |
| **Bosnia and Herz.** | UEFA | nfsbih.ba | `site:nfsbih.ba "World Cup"` |
| **Croatia** | UEFA | hns.family | `site:hns.family "World Cup"` |
| **Czechia** | UEFA | fotbal.cz | `site:fotbal.cz "World Cup"` |
| **England** | UEFA | englandfootball.com | `site:englandfootball.com "World Cup"` |
| **France** | UEFA | fff.fr | `site:fff.fr "Coupe du Monde"` |
| **Germany** | UEFA | dfb.de | `site:dfb.de "Weltmeisterschaft"` |
| **Netherlands** | UEFA | knvb.nl | `site:knvb.nl "World Cup"` |
| **Norway** | UEFA | fotball.no | `site:fotball.no "World Cup"` |
| **Portugal** | UEFA | fpf.pt | `site:fpf.pt "Mundial"` |
| **Scotland** | UEFA | scottishfa.co.uk | `site:scottishfa.co.uk "World Cup"` |
| **Spain** | UEFA | rfef.es | `site:rfef.es "Mundial"` |
| **Sweden** | UEFA | svenskfotboll.se | `site:svenskfotboll.se "World Cup"` |
| **Switzerland** | UEFA | football.ch | `site:football.ch "World Cup"` |
| **Türkiye** | UEFA | tff.org | `site:tff.org "World Cup"` |
| **Algeria** | CAF | faf.dz | `site:faf.dz "Coupe du Monde"` |
| **Cabo Verde** | CAF | fcf.cv | `site:fcf.cv "Mundial"` |
| **DR Congo** | CAF | fecofa.cd | `site:fecofa.cd "Coupe du Monde"` |
| **Egypt** | CAF | efa.com.eg | `site:efa.com.eg "World Cup"` |
| **Ghana** | CAF | ghanafa.org | `site:ghanafa.org "World Cup"` |
| **Côte d'Ivoire** | CAF | fifciv.com | `site:fifciv.com "Coupe du Monde"` |
| **Morocco** | CAF | frmf.ma | `site:frmf.ma "Coupe du Monde"` |
| **Senegal** | CAF | fsfoot.sn | `site:fsfoot.sn "Coupe du Monde"` |
| **South Africa** | CAF | safa.net | `site:safa.net "World Cup"` |
| **Tunisia** | CAF | ftf.org.tn | `site:ftf.org.tn "Coupe du Monde"` |
| **Australia** | AFC | footballaustralia.com.au | `site:footballaustralia.com.au "World Cup"` |
| **IR Iran** | AFC | ffiri.ir | `site:ffiri.ir "World Cup"` |
| **Iraq** | AFC | ifa.iq | `site:ifa.iq "World Cup"` |
| **Japan** | AFC | jfa.jp | `site:jfa.jp "World Cup"` |
| **Jordan** | AFC | jfa.jo | `site:jfa.jo "World Cup"` |
| **Qatar** | AFC | qfa.qa | `site:qfa.qa "World Cup"` |
| **Saudi Arabia** | AFC | saff.com.sa | `site:saff.com.sa "World Cup"` |
| **Korea Republic** | AFC | kfa.or.kr | `site:kfa.or.kr "World Cup"` |
| **Uzbekistan** | AFC | ufa.uz | `site:ufa.uz "World Cup"` |
| **New Zealand** | OFC | nzfootball.co.nz | `site:nzfootball.co.nz "World Cup"` |

#### G. Environmental & Venue Constraints
*   **Intent:** The user asks about match conditions, stadium details, or why a team might underperform.
*   **Prescribed Knowledge & Routing:**
    *   **Altitude Factor:** Mexico City (Estadio Azteca) is at 7,200+ feet (2,200m+); Guadalajara is at 5,100+ feet. Flag travel and oxygen recovery fatigue for teams playing matches here.
    *   **Surface Factor:** Venues like MetLife, AT&T Stadium, and Mercedes-Benz Stadium use artificial turf natively but are installing temporary real grass pitches for FIFA. Flag potential seams, slickness, or higher hamstring/groin strain risks.
    *   **Climate/Travel Factor:** Teams traveling from cool Pacific Northwest venues to high-humidity summer venues face massive thermodynamic shifts. Always factor rest days and flight distances into tactical analysis.

#### H. Tournament Progression & Tie-Breaker Logic
*   **Directive:** When calculating or discussing who advances to the Round of 32 from the 3rd-place pool, evaluate teams strictly by the official FIFA Group Stage tie-breakers in this exact order:
    1. Highest number of points obtained in all group matches.
    2. Superior goal difference in all group matches.
    3. Highest number of goals scored in all group matches.
    4. Fair play points (lowest number of yellow/red cards).
    5. Drawing of lots by the FIFA organizing committee.
*   **Action:** If a user asks "Can [Team X] still qualify?", explicitly map out their point potential and goal difference requirements compared to the other current 3rd-place teams.

#### I. Broadcast & Media Rights Routing
*   **Intent:** The user asks where to watch, stream, or listen to a live match.
*   **Routing Logic:**
    *   **United States:** English -> FOX, FS1, Fox Sports App. Spanish -> Telemundo, Universo, Peacock.
    *   **Canada:** CTV, TSN, RDS.
    *   **Mexico:** TelevisaUnivision, TV Azteca, ViX.
    *   **United Kingdom:** BBC (BBC iPlayer), ITV (STV/ITVX).
*   **Rule:** If the user's location is unknown, default to providing the US/Host broadcasters, but add: *"Please specify your country if you are tuning in from outside North America."*

#### J. Real-Time Anchor Rules
*   **Directive:** For any query regarding active injuries, lineup changes, training camp drama, or breaking news, you must append the current month and year to your search tool query.
*   **Search Formulation:** `[Player/Team Name] [Event] May 2026`
*   **Recency Filter:** Sort search tool results by "Past 24 Hours" or "Past Week" to avoid pulling speculative articles from prior years.
"""

def upload_skill_to_gcs(bucket_name: str, skill_content: str, destination_blob_name: str = "world_cup_2026_skill.md"):
    """Uploads a string literal to a Google Cloud Storage bucket.

    Args:
        bucket_name (str): The name of the target GCS bucket.
        skill_content (str): The markdown string content to upload.
        destination_blob_name (str): The filename to save it as in GCS.
    """

    if not bucket_name:
        print("❌ SKILL_GCS_BUCKET is not defined. Cannot upload.")
        return

    try:
        # Initialize the GCS client
        client = storage.Client()
        print(f"Connecting to bucket: {bucket_name}")
        bucket = client.bucket(bucket_name)

        # Create a new blob (file) within the bucket
        blob = bucket.blob(destination_blob_name)

        # Upload the string data directly
        blob.upload_from_string(skill_content)
        print(f"✅ Skill successfully uploaded to gs://{bucket_name}/{destination_blob_name}")

    except Exception as e:
        print(f"❌ Failed to upload to GCS: {e}")

# Execute the uploa
upload_skill_to_gcs(bucket_name=SKILL_GCS_BUCKET, skill_content=skill)

factory = AgentFactory(project_id=PROJECT_ID, endpoint=ENDPOINT, location=LOCATION)
harness = AgentHarness(factory)
harness.multi_turn_interact()