#!/usr/bin/env python
# COPYRIGHT 2025 Thomas Grothe
import argparse
import os
import sys
import json
import ai_api_module

#this is the CLI program used to interact with the GAI python module. 

default_model = "gemini-2.0-flash-exp"

def main():
    """Main function to parse arguments and execute commands"""
    parser = argparse.ArgumentParser(description="CLI for calling AI APIs with file attachment support")
    parser.add_argument(
        "--apikey", default=os.getenv("GEMINI_API_KEY"), help="API key for the service"
    )
    parser.add_argument("--modelid", default=default_model, help=f"ID of the model to use (default={default_model})")
    parser.add_argument("--list-models", action="store_true", help="list the available models")
    
    # Conversation management
    parser.add_argument("--new", "-n", action="store_true", help="Start a new conversation")
    parser.add_argument("--save-conversation", metavar="FILE", help="Save conversation to file")
    parser.add_argument("--load-conversation", metavar="FILE", help="Load conversation from file")
    parser.add_argument("--list-conversations", action="store_true", help="List saved conversations")
    parser.add_argument("--clear", action="store_true", help="Clear conversation history")
    parser.add_argument("--system-instruction", metavar="TEXT", help="Set system instruction for the conversation")
    
    # Query parameters
    parser.add_argument("--query", "-q", default="", help="Query to be sent to the assistant")
    parser.add_argument("--topp", type=float, default=0.9, help="Top P value for the model")
    parser.add_argument("--temperature", type=float, default=0.7, help="Temperature value for the model")
    parser.add_argument("--max-tokens", type=int, default=8192, help="Maximum output tokens")
    parser.add_argument("--stream", action="store_true", help="Enable streaming response")
    parser.add_argument("--no-stream", action="store_true", help="Disable streaming response")
    parser.add_argument("--output_mode", "-o", default='text', help="Output mode (text or json)")
    parser.add_argument("--count-tokens", action="store_true", help="Count tokens in query without sending")
    
    parser.add_argument("extra_query", nargs="*")

    args = parser.parse_args()
    
    # Get API key
    if args.apikey:
        ai_api_module.api_key = args.apikey
    else:
        ai_api_module.get_api_key()
        if not ai_api_module.api_key:
            print("Error: API key not provided. Set GEMINI_API_KEY environment variable.")
            sys.exit(1)
    
    # Handle system instruction
    if args.system_instruction:
        ai_api_module.set_system_instruction(args.system_instruction)
        print(f"System instruction set")
    
    # Handle conversation management
    if args.new:
        ai_api_module.clear_conversation()
        print("Started new conversation")
    
    if args.load_conversation:
        if ai_api_module.load_conversation(args.load_conversation):
            print(f"Loaded conversation from {args.load_conversation}")
        else:
            print(f"Failed to load conversation from {args.load_conversation}")
            sys.exit(1)
    
    if args.list_conversations:
        success, result = ai_api_module.list_saved_conversations()
        if success:
            if not result:
                print("No saved conversations found")
            else:
                print("Saved conversations:")
                for conv in result:
                    print(f"  {conv['filename']}")
                    print(f"    Saved: {conv['saved_at']}")
                    print(f"    Messages: {conv['message_count']}")
                    print()
        else:
            print(f"Error: {result}")
        return
    
    if args.clear:
        ai_api_module.clear_conversation()
        print("Conversation history cleared")
        return
    
    # List models
    if args.list_models:
        success, result = ai_api_module.list_models(args.output_mode)
        if success:
            for model in result:
                if args.output_mode == 'json':
                    print(json.dumps(model, indent=2))
                else:
                    print(f"Model: {model['display_name']}")
                    print(f"  Name: {model['name']}")
                    print(f"  Description: {model.get('description', 'N/A')}")
                    print(f"  Input tokens: {model.get('input_token_limit', 'N/A')}")
                    print(f"  Output tokens: {model.get('output_token_limit', 'N/A')}")
                    print()
        else:
            print(f'Error: {result}')
        return

    # Process query
    if args.query or args.extra_query:
        query = args.query
        if args.extra_query:
            if query:
                query += " " + " ".join(args.extra_query)
            else:
                query = " ".join(args.extra_query)
        
        # Count tokens if requested
        if args.count_tokens:
            success, count = ai_api_module.count_tokens(args.modelid, query)
            if success:
                print(f"Token count: {count}")
            else:
                print(f"Error counting tokens: {count}")
            return
        
        # Determine streaming
        use_stream = args.stream and not args.no_stream
        
        response = ai_api_module.query( 
            args.modelid, 
            query,
            args.topp,
            args.temperature, 
            use_stream,
            args.max_tokens,
        )
        
        if response and not use_stream:
            print(response)
        
        # Auto-save conversation if requested
        if args.save_conversation:
            filepath = ai_api_module.save_conversation(args.save_conversation)
            print(f"\nConversation saved to {filepath}")
    else:
        parser.print_help()

if __name__ == "__main__":
    main()