import logging
import argparse
from typing import List
from models.ollama_client import OllamaClient
from benchmarks.translation import TranslationBenchmark
from benchmarks.summarization import SummarizationBenchmark
from benchmarks.codegen import CodegenBenchmark
from utils.report import ReportGenerator

def setup_logging(verbose: bool = False):
    """Настройка логирования."""
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler()
        ]
    )

def run_benchmarks(ollama_client: OllamaClient, model_name: str, benchmarks: List[str], context_size: int = None) -> List[dict]:
    """
    Запуск выбранных бенчмарков.

    Args:
        ollama_client: Клиент для работы с Ollama
        model_name: Название модели
        benchmarks: Список имен бенчмарков для запуска
        context_size: Размер контекста для модели

    Returns:
        Список результатов бенчмарков
    """
    benchmark_classes = {
        'translation': TranslationBenchmark,
        'summarization': SummarizationBenchmark,
        'codegen': CodegenBenchmark
    }

    results = []

    for benchmark_name in benchmarks:
        if benchmark_name not in benchmark_classes:
            logging.warning(f"Unknown benchmark: {benchmark_name}")
            continue

        logging.info(f"Running {benchmark_name} benchmark...")
        benchmark = benchmark_classes[benchmark_name]()
        result = benchmark.run(ollama_client, model_name, context_size)
        results.append(result)

    return results

def main():
    """Основная функция запуска."""
    parser = argparse.ArgumentParser(description='LLM Benchmarking Tool')
    parser.add_argument('-m', '--model', required=True, help='Название модели для тестирования')
    parser.add_argument('-u', '--ollama-url', default='http://localhost:11434', help='URL подключения к Ollama серверу')
    parser.add_argument('-c', '--context-size', type=int, default=32000, help='Размер контекста для модели (по умолчанию 32000)')
    parser.add_argument('-b', '--benchmarks', nargs='+', default=['translation', 'summarization', 'codegen'],
                       help='Список бенчмарков для выполнения (translation, summarization, codegen)')
    parser.add_argument('-o', '--output', default='results', help='Директория для сохранения результатов')
    parser.add_argument('-v', '--verbose', action='store_true', help='Подробный режим вывода')

    args = parser.parse_args()

    # Настройка логирования
    setup_logging(args.verbose)

    logging.info(f"Starting benchmarking for model: {args.model}")
    logging.info(f"Ollama URL: {args.ollama_url}")
    logging.info(f"Benchmarks to run: {', '.join(args.benchmarks)}")
    logging.info(f"Context size: {args.context_size}")

    # Инициализация клиента
    ollama_client = OllamaClient(args.ollama_url)
    try:
        # Запуск бенчмарков
        results = run_benchmarks(ollama_client, args.model, args.benchmarks, args.context_size)

        # Генерация отчетов
        report_generator = ReportGenerator()

        for result in results:
            report_generator.generate_benchmark_report(result, args.output, args.model)

        if len(results) > 1:
            report_generator.generate_summary_report(results, args.output, args.model, args.ollama_url)

        logging.info("Benchmarking completed successfully!")
    except Exception as e:
        logging.error(f"Error during benchmarking: {e}")
        return 1

    return 0

if __name__ == '__main__':
    exit(main())
