"""
Fine-tune Agent Script
Usage: python scripts/fine_tune_agent.py --agent nutrition --min-rating 4.0
"""

import argparse
from fine_tuning import get_data_collector, fine_tune_agent


def main():
    parser = argparse.ArgumentParser(description='Fine-tune a healthcare agent')
    parser.add_argument('--agent', required=True, 
                       choices=['nutrition', 'exercise', 'symptom', 'mental_health', 'general_health'],
                       help='Agent to fine-tune')
    parser.add_argument('--min-rating', type=float, default=None,
                       help='Minimum quality rating (1-5) to include conversations')
    parser.add_argument('--model', default='gpt-4o-mini-2024-07-18',
                       help='Base model to fine-tune')
    parser.add_argument('--suffix', default=None,
                       help='Suffix for fine-tuned model name')
    parser.add_argument('--no-wait', action='store_true',
                       help='Don\'t wait for fine-tuning to complete')
    
    args = parser.parse_args()
    
    # Get data collector
    collector = get_data_collector()
    
    # Check conversation count
    counts = collector.get_conversation_count(f"{args.agent}_agent")
    agent_key = args.agent
    
    if agent_key not in counts or counts[agent_key] == 0:
        print(f"❌ No conversations found for {args.agent} agent")
        print(f"   Start using the chatbot to collect training data")
        return
    
    print(f"📊 Found {counts[agent_key]} conversations for {args.agent} agent")
    
    # Export training data
    print(f"\n📤 Exporting training data...")
    training_file = collector.export_for_openai_finetuning(
        agent_name=f"{args.agent}_agent",
        min_quality_rating=args.min_rating
    )
    
    # Start fine-tuning
    print(f"\n🚀 Starting fine-tuning job...")
    result = fine_tune_agent(
        agent_name=args.agent,
        training_file=training_file,
        model=args.model,
        suffix=args.suffix,
        wait_for_completion=not args.no_wait
    )
    
    if args.no_wait:
        print(f"\n✅ Fine-tuning job started: {result}")
        print(f"   Check status with: python scripts/check_finetuning_status.py --job-id {result}")
    else:
        print(f"\n✅ Fine-tuning completed!")
        print(f"   Model ID: {result}")
        print(f"\n💡 To use this model, update your agent configuration:")
        print(f"   MODEL = '{result}'")


if __name__ == '__main__':
    main()