plugin.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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
#
# http://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.

"""Game Character AI G-Assist plugin - Main entry point."""

import json
import logging
import os
from ctypes import byref, windll, wintypes
from typing import Optional

try:
    from groq import Groq
except ImportError:
    Groq = None

from ai_backends import CharacterAnalyzer, LMStudioClient, SupabaseClient
from config_manager import ConfigManager
from models import Response
from window_manager import WindowManager

# Global variables
config_manager = None
window_manager = None
character_analyzer = None
character_context = {}  # Store character information

# Setup logging
LOG_FILE = os.path.join(os.environ.get("USERPROFILE", "."), "game-character-ai.log")
logging.basicConfig(
    filename=LOG_FILE,
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)


def main():
    """Main entry point.

    Sits in a loop listening to a pipe, waiting for commands to be issued. After
    receiving the command, it is processed and the result returned. The loop
    continues until the "shutdown" command is issued.
    """
    # Initialize the plugin on startup
    execute_initialize_command()

    while True:
        command = read_command()
        if command is None:
            logging.error("Error reading command")
            continue

        tool_calls = command.get("tool_calls", [])
        context = command.get("messages", [])
        system_info = command.get("system_info", {})

        for tool_call in tool_calls:
            func = tool_call.get("func")
            params = tool_call.get("properties", {})

            if func == "initialize":
                response = execute_initialize_command()
            elif func == "screenshot":
                response = execute_analyze_character_command(
                    params, context, system_info
                )
            elif func == "clear_context":
                response = execute_clear_context_command(params)
            elif func == "shutdown":
                response = execute_shutdown_command()
                write_response(response)
                return
            else:
                response = generate_failure_response(f"Unknown function call: {func}")

            write_response(response)


def read_command() -> dict | None:
    """Reads a command from the communication pipe.

    Returns:
        Command details if the input was proper JSON; `None` otherwise
    """
    retval = ""
    try:
        STD_INPUT_HANDLE = -10
        pipe = windll.kernel32.GetStdHandle(STD_INPUT_HANDLE)

        chunks = []
        while True:
            BUFFER_SIZE = 4096
            message_bytes = wintypes.DWORD()
            buffer = bytes(BUFFER_SIZE)
            success = windll.kernel32.ReadFile(
                pipe, buffer, BUFFER_SIZE, byref(message_bytes), None
            )

            if not success:
                logging.error("Error reading from command pipe")
                return None

            chunk = buffer.decode("utf-8")[: message_bytes.value]
            chunks.append(chunk)

            if message_bytes.value < BUFFER_SIZE:
                break

        retval = "".join(chunks).strip()

        # Skip empty or whitespace-only messages
        if not retval:
            return None

        return json.loads(retval)

    except json.JSONDecodeError:
        logging.error(f"Received invalid JSON: '{retval}' (length: {len(retval)})")
        return None
    except Exception as e:
        logging.error(f"Exception in read_command(): {str(e)}")
        return None


def write_response(response: Response) -> None:
    """Writes a response to the communication pipe.

    Args:
        response: Dictionary containing return value(s)
    """
    try:
        STD_OUTPUT_HANDLE = -11
        pipe = windll.kernel32.GetStdHandle(STD_OUTPUT_HANDLE)

        json_message = json.dumps(response) + "<<END>>"
        message_bytes = json_message.encode("utf-8")
        message_len = len(message_bytes)

        bytes_written = wintypes.DWORD()
        windll.kernel32.WriteFile(pipe, message_bytes, message_len, bytes_written, None)

    except Exception:
        logging.error("Unknown exception caught.")
        pass


def generate_failure_response(message: Optional[str] = None) -> Response:
    """Generates a response indicating failure.

    Args:
        message: String to be returned in the response (optional)

    Returns:
        A failure response with the attached message
    """
    response: Response = {"success": False}
    if message:
        response["message"] = message
    return response


def generate_success_response(message: Optional[str] = None) -> Response:
    """Generates a response indicating success.

    Args:
        message: String to be returned in the response (optional)

    Returns:
        A success response with the attached message
    """
    response: Response = {"success": True}
    if message:
        response["message"] = message
    return response


def generate_message_response(message: str) -> Response:
    """Generates a message response.

    Args:
        message: String to be returned to the driver

    Returns:
        A message response dictionary
    """
    response: Response = {"message": message}
    return response


def execute_initialize_command() -> dict:
    """Initialize the plugin and external service clients.

    Returns:
        Success response
    """
    global config_manager, window_manager, character_analyzer

    try:
        # Initialize configuration manager
        config_manager = ConfigManager()
        config_manager.load_config()

        # Initialize Groq client if API key is present
        groq_client = None
        groq_api_key = config_manager.get_groq_api_key()
        if groq_api_key and Groq:
            groq_client = Groq(api_key=groq_api_key)

        # Initialize window manager
        window_manager = WindowManager(groq_client=groq_client)

        # Initialize AI backends
        supabase_client = None
        supabase_endpoint = config_manager.get_supabase_endpoint()
        if supabase_endpoint:
            supabase_client = SupabaseClient(
                endpoint_url=supabase_endpoint,
                auth_header=config_manager.get_supabase_auth_header(),
            )

        lm_studio_client = None
        lm_studio_endpoint = config_manager.get_lm_studio_endpoint()
        if lm_studio_endpoint:
            lm_studio_client = LMStudioClient(
                endpoint=lm_studio_endpoint,
                api_key=config_manager.get_lm_studio_api_key(),
            )

        # Initialize character analyzer
        character_analyzer = CharacterAnalyzer(
            groq_client=groq_client,
            supabase_client=supabase_client,
            lm_studio_client=lm_studio_client,
        )

        return generate_success_response()

    except Exception as e:
        logging.error(f"Initialization failed: {str(e)}")
        # Return success to allow the plugin to continue
        return generate_success_response()


def execute_shutdown_command() -> dict:
    """Cleanup resources.

    Returns:
        Success response
    """
    return generate_success_response()


def execute_analyze_character_command(
    params: Optional[dict] = None,
    context: Optional[list] = None,
    system_info: Optional[dict] = None,
) -> dict:
    """Capture screenshot, identify game character, and optionally answer a question.

    Args:
        params: Parameters containing optional question
        context: Conversation context
        system_info: System information including current game

    Returns:
        Success response with character analysis
    """
    global character_context, window_manager, character_analyzer

    if not window_manager or not character_analyzer:
        return generate_failure_response("Plugin not properly initialized")

    try:
        # Extract current game from system_info
        current_game = "Unknown"
        if system_info:
            if isinstance(system_info, dict):
                current_game = system_info.get(
                    "current_game", system_info.get("game", "Unknown")
                )
            elif isinstance(system_info, str):
                current_game = system_info

        # Capture screenshot
        write_response(generate_message_response("📸 Taking screenshot..."))
        screenshot_base64 = window_manager.capture_screenshot(current_game)

        # Tell user where screenshot was saved
        screenshot_path = os.path.join(
            os.environ.get("USERPROFILE", "."), "gamecharacterai_last_screenshot.png"
        )
        write_response(
            generate_message_response(f"💾 Screenshot saved to: {screenshot_path}")
        )

        # Query Supabase character database if available
        supabase_results = []
        if character_analyzer.supabase_client:
            write_response(
                generate_message_response("🔍 Searching character database...")
            )
            supabase_results = character_analyzer.supabase_client.query_character_db(
                screenshot_base64, current_game, limit=3
            )

        # Get the user's question/prompt from the conversation context
        user_prompt = ""
        if context and len(context) > 0:
            user_prompt = context[-1].get("content", "")
            logging.info(f"User asked: '{user_prompt}'")

        # Get optional specific question from params (legacy support)
        question = params.get("question", "") if params else ""

        # Use the context prompt if available, otherwise fall back to question param
        actual_user_request = user_prompt if user_prompt else question
        logging.info(f"Final user request: '{actual_user_request}'")

        # Use character analyzer to get response
        write_response(generate_message_response("🤖 Analyzing character..."))

        try:
            character_info = character_analyzer.analyze_character(
                image_base64=screenshot_base64,
                current_game=current_game,
                user_request=actual_user_request,
                supabase_results=supabase_results,
            )

            # Ensure response is reasonably sized
            if isinstance(character_info, str) and len(character_info) > 250:
                character_info = character_info[:250] + "..."

        except Exception as api_error:
            logging.error(f"Character analysis failed: {str(api_error)}")
            character_info = "Character analysis failed. Please try again."

        # Store character context for potential follow-up questions
        character_context = {
            "game": current_game,
            "character_info": character_info,
            "screenshot": screenshot_base64,
            "last_question": question,
        }

        # Add a newline before the AI response for better readability
        write_response(generate_message_response(f"\n{character_info}"))
        return generate_success_response()

    except Exception as e:
        logging.error(f"Character analysis failed: {str(e)}")
        return generate_failure_response(f"Failed to analyze character: {str(e)}")


def execute_clear_context_command(params: Optional[dict] = None) -> dict:
    """Clear the stored character context.

    Returns:
        Success response
    """
    global character_context
    character_context = {}
    return generate_success_response(
        "Character context cleared. Ready to identify a new character."
    )


if __name__ == "__main__":
    main()