import argparse
import json
from pathlib import Path
from datetime import datetime

from components import GenerationConfig, GenerationPipeline


def main():
               
    parser = argparse.ArgumentParser(description='Standalone Generation Runner')

          
    parser.add_argument('--retrieval_file', help='Retrieval results file path')
    parser.add_argument('--output_file', help='Output file path (optional, auto-generated)')

          
    parser.add_argument('--api_key', default='{your api_key}')
    parser.add_argument('--base_url', default='{your base_url}')
    parser.add_argument('--llm_model', default='gpt-4o', help='Select generation model')
    parser.add_argument('--temperature', type=float, default=0.1)
    parser.add_argument('--concurrent_requests', type=int, default=3, help='Number of concurrent requests')

          
    parser.add_argument('--verbose', action='store_true', default=True, help='Verbose output')
    parser.add_argument('--list_models', action='store_true', help='List all available models')

    args = parser.parse_args()

                       
    if args.list_models:
        models = ['gpt-4o', 'gemini-2.5-flash', 'deepseek-chat', 'llama-4-scout', 'llama-3.1-8b-instruct', 'qwen3-14b']
        for i, model in enumerate(models, 1):
            print(f"{i}. {model}")
        return

                   
    if not args.retrieval_file:
        print("Please provide a retrieval file path")
        return

              
    retrieval_file = Path(args.retrieval_file)
    if not retrieval_file.exists():
        print(f"Retrieval file {retrieval_file} does not exist")
        return

            
    config = GenerationConfig(
        api_key=args.api_key,
        base_url=args.base_url,
        llm_model=args.llm_model,
        temperature=args.temperature,
        concurrent_requests=args.concurrent_requests,
        verbose=args.verbose
    )


            
    retrieval_results = []
    try:
        with open(retrieval_file, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                if line.strip():
                    try:
                        data = json.loads(line.strip())
                        retrieval_results.append(data)
                    except json.JSONDecodeError as e:
                        continue


    except Exception as e:
        print(f"Error loading retrieval results: {e}")
        return

    if not retrieval_results:
        print("No retrieval results found")
        return

                   
    pipeline = GenerationPipeline(config)

            
    if args.output_file:
        output_file = Path(args.output_file)
    else:
                 
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        model_name = args.llm_model.replace('-', '_').replace('/', '_')
        temp_str = str(args.temperature).replace('.', '')
        filename = f"generation_{retrieval_file.stem}_{model_name}_t{temp_str}_c{args.concurrent_requests}_{timestamp}.jsonl"

        output_dir = Path(config.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        output_file = output_dir / filename


          
    results = pipeline.batch_generate(retrieval_results, str(output_file))


            
    success_count = len([r for r in results if 'error' not in r])
    error_count = len(results) - success_count

    if error_count > 0:
        print(f"Error generating for {error_count} items")  
            
    if success_count > 0 and args.verbose:
        sample_result = next(r for r in results if 'error' not in r)
        print(f"Successfully generated for {success_count} items")
        print(f"Sample result: {sample_result}")
    else:
        print("No successful generations")


if __name__ == '__main__':
    main()
