"""
Django management команда для поиска мест в коде, где делаются запросы к БД.
Находит:
1. Queryset'ы в views.py
2. SerializerMethodField с запросами к БД
3. Модели с запросами в методах/properties
4. Services с запросами
"""
import ast
import os
from pathlib import Path
from django.core.management.base import BaseCommand


class DatabaseQueryFinder(ast.NodeVisitor):
    """AST visitor для поиска запросов к БД."""
    
    def __init__(self, file_path):
        self.file_path = file_path
        self.queries_found = []
        self.current_class = None
        self.current_method = None
    
    def visit_ClassDef(self, node):
        """Посещает классы."""
        old_class = self.current_class
        self.current_class = node.name
        self.generic_visit(node)
        self.current_class = old_class
    
    def visit_FunctionDef(self, node):
        """Посещает функции/методы."""
        old_method = self.current_method
        self.current_method = node.name
        self.generic_visit(node)
        self.current_method = old_method
    
    def visit_Call(self, node):
        """Посещает вызовы функций."""
        # Ищем queryset методы
        queryset_methods = [
            'filter', 'get', 'create', 'update', 'delete',
            'exclude', 'annotate', 'aggregate', 'count',
            'exists', 'values', 'values_list', 'select_related',
            'prefetch_related', 'only', 'defer', 'first', 'last',
            'all', 'order_by', 'distinct'
        ]
        
        # Проверяем вызовы типа obj.objects.filter(), obj.filter()
        if isinstance(node.func, ast.Attribute):
            attr_name = node.func.attr
            if attr_name in queryset_methods:
                # Определяем контекст
                context = self._get_context()
                
                self.queries_found.append({
                    'line': node.lineno,
                    'method': attr_name,
                    'context': context,
                    'code': self._get_code_snippet(node),
                })
        
        # Ищем вызовы моделей напрямую (например, User.objects.get())
        if isinstance(node.func, ast.Attribute):
            if node.func.attr == 'objects':
                context = self._get_context()
                self.queries_found.append({
                    'line': node.lineno,
                    'method': 'objects',
                    'context': context,
                    'code': self._get_code_snippet(node),
                })
        
        self.generic_visit(node)
    
    def _get_context(self):
        """Получает контекст (класс.метод)."""
        if self.current_class and self.current_method:
            return f"{self.current_class}.{self.current_method}"
        elif self.current_method:
            return self.current_method
        elif self.current_class:
            return f"{self.current_class}."
        return "global"
    
    def _get_code_snippet(self, node, max_length=100):
        """Получает фрагмент кода (упрощенная версия)."""
        # В реальной версии нужно было бы использовать tokenize или другой метод
        return f"Line {node.lineno}"


class Command(BaseCommand):
    help = 'Находит места в коде где делаются запросы к БД'
    
    def add_arguments(self, parser):
        parser.add_argument(
            '--app',
            type=str,
            help='Имя приложения для анализа (например: homework, chat)',
        )
        parser.add_argument(
            '--file',
            type=str,
            help='Конкретный файл для анализа',
        )
        parser.add_argument(
            '--serializers',
            action='store_true',
            help='Проверить только serializers',
        )
        parser.add_argument(
            '--views',
            action='store_true',
            help='Проверить только views',
        )
    
    def handle(self, *args, **options):
        app_name = options.get('app')
        file_path = options.get('file')
        check_serializers = options.get('serializers')
        check_views = options.get('views')
        
        self.stdout.write(self.style.SUCCESS('=' * 70))
        self.stdout.write(self.style.SUCCESS('🔍 Поиск запросов к БД в коде'))
        self.stdout.write(self.style.SUCCESS('=' * 70))
        self.stdout.write('')
        
        base_dir = Path(__file__).resolve().parent.parent.parent.parent
        apps_dir = base_dir / 'apps'
        
        files_to_check = []
        
        if file_path:
            # Конкретный файл
            file_path = Path(file_path)
            if file_path.exists():
                files_to_check.append(file_path)
            else:
                self.stdout.write(self.style.ERROR(f'Файл не найден: {file_path}'))
                return
        elif app_name:
            # Все файлы приложения
            app_dir = apps_dir / app_name
            if not app_dir.exists():
                self.stdout.write(self.style.ERROR(f'Приложение не найдено: {app_name}'))
                return
            
            if check_serializers:
                files_to_check.append(app_dir / 'serializers.py')
            elif check_views:
                files_to_check.append(app_dir / 'views.py')
            else:
                files_to_check.extend([
                    app_dir / 'views.py',
                    app_dir / 'serializers.py',
                ])
        else:
            # Все приложения
            for app_dir in apps_dir.iterdir():
                if app_dir.is_dir() and (app_dir / '__init__.py').exists():
                    if check_serializers:
                        ser_file = app_dir / 'serializers.py'
                        if ser_file.exists():
                            files_to_check.append(ser_file)
                    elif check_views:
                        views_file = app_dir / 'views.py'
                        if views_file.exists():
                            files_to_check.append(views_file)
                    else:
                        views_file = app_dir / 'views.py'
                        ser_file = app_dir / 'serializers.py'
                        if views_file.exists():
                            files_to_check.append(views_file)
                        if ser_file.exists():
                            files_to_check.append(ser_file)
        
        total_queries = 0
        files_with_issues = []
        
        for file_path in files_to_check:
            if not file_path.exists():
                continue
            
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                    tree = ast.parse(content, filename=str(file_path))
                
                finder = DatabaseQueryFinder(file_path)
                finder.visit(tree)
                
                if finder.queries_found:
                    relative_path = file_path.relative_to(base_dir)
                    self.stdout.write(self.style.WARNING(f'\n📄 {relative_path}'))
                    self.stdout.write('-' * 70)
                    
                    # Группируем по контексту
                    by_context = {}
                    for query in finder.queries_found:
                        context = query['context']
                        if context not in by_context:
                            by_context[context] = []
                        by_context[context].append(query)
                    
                    for context, queries in by_context.items():
                        self.stdout.write(f'   📍 {context}: {len(queries)} запросов')
                        for query in queries[:5]:  # Показываем первые 5
                            self.stdout.write(f'      Line {query["line"]}: {query["method"]}')
                        if len(queries) > 5:
                            self.stdout.write(f'      ... и еще {len(queries) - 5}')
                    
                    total_queries += len(finder.queries_found)
                    files_with_issues.append((file_path, finder.queries_found))
            
            except SyntaxError as e:
                self.stdout.write(self.style.ERROR(f'Ошибка парсинга {file_path}: {e}'))
            except Exception as e:
                self.stdout.write(self.style.ERROR(f'Ошибка при обработке {file_path}: {e}'))
        
        self.stdout.write('')
        self.stdout.write(self.style.SUCCESS('=' * 70))
        self.stdout.write(f'Всего найдено: {total_queries} потенциальных запросов в {len(files_with_issues)} файлах')
        self.stdout.write('')
        self.stdout.write('💡 Рекомендации:')
        self.stdout.write('   1. Проверьте SerializerMethodField - часто источник N+1 проблем')
        self.stdout.write('   2. Используйте select_related() и prefetch_related() в queryset')
        self.stdout.write('   3. Используйте aggregate() вместо множественных count()')
        self.stdout.write('   4. Проверьте наличие индексов в БД для часто используемых полей')

