import argparse
import os
import sys
from typing import Any, Dict

from dotenv import load_dotenv
from google import genai
from google.genai import types

from functions.call_function import available_functions, call_function


def main() -> None:
    parser = argparse.ArgumentParser(description="AI Code Assistant")
    parser.add_argument("user_prompt", type=str, help="Prompt to send to Gemini")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
    args = parser.parse_args()

    load_dotenv()
    api_key = os.environ.get("GEMINI_API_KEY")

    if not api_key:
        raise RuntimeError("GEMINI_API_KEY environment variable not set")

    client = genai.Client(api_key=api_key)

    messages: list[types.Content] = [
        types.Content(role="user", parts=[types.Part(text=args.user_prompt)])
    ]

    if args.verbose:
        print(f"User prompt: {args.user_prompt}\n")

    generate_content(client, messages, args.verbose)


def generate_content(
    client: genai.Client, messages: list[types.Content], verbose: bool
) -> None:
    system_instructions = """
    You are a helpful AI coding agent.

    When a user asks a question or makes a request, make a function call plan. You can perform the following operations:

    - List files and directories
    - Read file contents
    - Execute Python files with optional arguments
    - Write or overwrite files

    All paths you provide should be relative to the working directory. You do not need to specify the working directory in your function calls as it is automatically injected for security reasons.
    """

    for _ in range(20):
        model_response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=messages,
            config=types.GenerateContentConfig(
                system_instruction=system_instructions, tools=[available_functions]
            ),
        )

        if not model_response.usage_metadata:
            raise RuntimeError("Gemini API response appears to be malformed")

        if verbose:
            print("Prompt tokens:", model_response.usage_metadata.prompt_token_count)
            print(
                "Response tokens:", model_response.usage_metadata.candidates_token_count
            )

        if model_response.candidates != None and len(model_response.candidates) != 0:
            for candidate in model_response.candidates:
                if candidate.content != None:
                    messages.append(candidate.content)

        if not model_response.function_calls:
            print("Final Response:")
            print(model_response.text)
            return

        function_responses: list[types.Part] = []

        for function_call in model_response.function_calls:
            result = call_function(function_call, verbose)

            if (
                not result.parts
                or not result.parts[0].function_response
                or not result.parts[0].function_response.response
            ):
                raise RuntimeError(f"Empty function response for {function_call.name}")

            if verbose:
                print(f"-> {result.parts[0].function_response.response}")

            function_responses.append(result.parts[0])

        messages.append(types.Content(role="user", parts=function_responses))
        continue

    print("Model took more than the max iterations to produce final response")
    sys.exit(1)


if __name__ == "__main__":
    main()