#!/usr/bin/env python
"""
Script to update the max_tokens parameter for OpenRouter models in the configuration.
"""

import json
import logging
import os
from config.config import get_config

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def update_openrouter_max_tokens(model_name="openrouter-claude-3.7-sonnet", new_max_tokens=8000):
    """
    Update the max_tokens parameter for an OpenRouter model in the configuration.
    
    Args:
        model_name: Name of the OpenRouter model to update
        new_max_tokens: New value for max_tokens parameter
    """
    logger.info(f"Updating max_tokens for {model_name} to {new_max_tokens}...")
    
    # Get the config instance
    config = get_config()
    
    # Check if the model exists in the configuration
    if 'models' not in config.config_data:
        logger.error("No models section found in configuration")
        return False
    
    if model_name not in config.config_data['models']:
        logger.error(f"Model {model_name} not found in configuration")
        return False
    
    # Get the current model configuration
    model_config = config.config_data['models'][model_name]
    
    # Print current configuration
    logger.info(f"Current configuration for {model_name}:")
    logger.info(json.dumps(model_config, indent=2))
    
    # Update the max_tokens parameter
    old_max_tokens = model_config.get('max_tokens', 2048)
    model_config['max_tokens'] = new_max_tokens
    
    # Update the configuration
    config.config_data['models'][model_name] = model_config
    
    # Save the configuration (in-memory only, as we can't modify the file directly)
    logger.info(f"Updated max_tokens for {model_name} from {old_max_tokens} to {new_max_tokens}")
    logger.info(f"New configuration for {model_name}:")
    logger.info(json.dumps(model_config, indent=2))
    
    logger.info("Configuration updated in memory. The next time you run a report, it will use the new max_tokens value.")
    logger.info("Note: This change is temporary and will be reset when the application restarts.")
    logger.info("To make the change permanent, you need to update the config.yaml file directly.")
    
    return True

def main():
    """Main function."""
    # Update max_tokens for Claude 3.7 Sonnet
    update_openrouter_max_tokens("openrouter-claude-3.7-sonnet", 8000)
    
    # You can also update other OpenRouter models if needed
    # update_openrouter_max_tokens("openrouter-mixtral", 8000)

if __name__ == "__main__":
    main()