"""
Storage module для PyGuardian
Управление SQLite базой данных для хранения IP-адресов, попыток атак и банов
"""

import asyncio
import sqlite3
import aiosqlite
import ipaddress
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Tuple
import logging

logger = logging.getLogger(__name__)


class Storage:
    """Асинхронный класс для работы с SQLite базой данных"""
    
    def __init__(self, db_path: str):
        self.db_path = db_path
        self._connection = None
    
    async def init_database(self) -> None:
        """Инициализация базы данных и создание таблиц"""
        async with aiosqlite.connect(self.db_path) as db:
            # Таблица для хранения попыток атак
            await db.execute("""
                CREATE TABLE IF NOT EXISTS attack_attempts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    ip_address TEXT NOT NULL,
                    username TEXT,
                    attack_type TEXT NOT NULL,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    log_line TEXT,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(ip_address),
                    INDEX(timestamp)
                )
            """)
            
            # Таблица для хранения забаненных IP
            await db.execute("""
                CREATE TABLE IF NOT EXISTS banned_ips (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    ip_address TEXT UNIQUE NOT NULL,
                    ban_reason TEXT NOT NULL,
                    banned_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    unban_at DATETIME,
                    is_active BOOLEAN DEFAULT 1,
                    manual_ban BOOLEAN DEFAULT 0,
                    attempts_count INTEGER DEFAULT 0,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(ip_address),
                    INDEX(unban_at),
                    INDEX(is_active)
                )
            """)
            
            # Таблица для успешных входов
            await db.execute("""
                CREATE TABLE IF NOT EXISTS successful_logins (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    ip_address TEXT NOT NULL,
                    username TEXT NOT NULL,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    session_info TEXT,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(ip_address),
                    INDEX(username),
                    INDEX(timestamp)
                )
            """)
            
            # Таблица для статистики
            await db.execute("""
                CREATE TABLE IF NOT EXISTS daily_stats (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    date DATE UNIQUE NOT NULL,
                    total_attempts INTEGER DEFAULT 0,
                    unique_ips INTEGER DEFAULT 0,
                    banned_count INTEGER DEFAULT 0,
                    successful_logins INTEGER DEFAULT 0,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(date)
                )
            """)
            
            await db.commit()
            logger.info("База данных инициализирована успешно")
    
    async def add_attack_attempt(self, ip: str, username: str, 
                               attack_type: str, log_line: str, 
                               timestamp: Optional[datetime] = None) -> None:
        """Добавить попытку атаки в базу данных"""
        if timestamp is None:
            timestamp = datetime.now()
            
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                INSERT INTO attack_attempts 
                (ip_address, username, attack_type, timestamp, log_line)
                VALUES (?, ?, ?, ?, ?)
            """, (ip, username, attack_type, timestamp, log_line))
            await db.commit()
    
    async def get_attack_count_for_ip(self, ip: str, time_window: int) -> int:
        """Получить количество попыток атак от IP за указанный период"""
        since_time = datetime.now() - timedelta(seconds=time_window)
        
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT COUNT(*) FROM attack_attempts 
                WHERE ip_address = ? AND timestamp > ?
            """, (ip, since_time)) as cursor:
                result = await cursor.fetchone()
                return result[0] if result else 0
    
    async def ban_ip(self, ip: str, reason: str, unban_time: int, 
                    manual: bool = False, attempts_count: int = 0) -> bool:
        """Забанить IP адрес"""
        unban_at = datetime.now() + timedelta(seconds=unban_time)
        
        async with aiosqlite.connect(self.db_path) as db:
            try:
                await db.execute("""
                    INSERT OR REPLACE INTO banned_ips 
                    (ip_address, ban_reason, unban_at, manual_ban, attempts_count)
                    VALUES (?, ?, ?, ?, ?)
                """, (ip, reason, unban_at, manual, attempts_count))
                await db.commit()
                logger.info(f"IP {ip} забанен. Причина: {reason}")
                return True
            except Exception as e:
                logger.error(f"Ошибка при бане IP {ip}: {e}")
                return False
    
    async def unban_ip(self, ip: str) -> bool:
        """Разбанить IP адрес"""
        async with aiosqlite.connect(self.db_path) as db:
            try:
                await db.execute("""
                    UPDATE banned_ips 
                    SET is_active = 0 
                    WHERE ip_address = ? AND is_active = 1
                """, (ip,))
                await db.commit()
                logger.info(f"IP {ip} разбанен")
                return True
            except Exception as e:
                logger.error(f"Ошибка при разбане IP {ip}: {e}")
                return False
    
    async def is_ip_banned(self, ip: str) -> bool:
        """Проверить, забанен ли IP"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT id FROM banned_ips 
                WHERE ip_address = ? AND is_active = 1 
                AND (unban_at IS NULL OR unban_at > datetime('now'))
            """, (ip,)) as cursor:
                result = await cursor.fetchone()
                return result is not None
    
    async def get_expired_bans(self) -> List[str]:
        """Получить список IP с истекшим временем бана"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT ip_address FROM banned_ips 
                WHERE is_active = 1 AND unban_at <= datetime('now')
                AND manual_ban = 0
            """) as cursor:
                results = await cursor.fetchall()
                return [row[0] for row in results]
    
    async def get_banned_ips(self) -> List[Dict]:
        """Получить список всех забаненных IP"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT ip_address, ban_reason, banned_at, unban_at, 
                       manual_ban, attempts_count
                FROM banned_ips 
                WHERE is_active = 1
                ORDER BY banned_at DESC
            """) as cursor:
                results = await cursor.fetchall()
                
                banned_list = []
                for row in results:
                    banned_list.append({
                        'ip': row[0],
                        'reason': row[1],
                        'banned_at': row[2],
                        'unban_at': row[3],
                        'manual': bool(row[4]),
                        'attempts': row[5]
                    })
                return banned_list
    
    async def get_top_attackers(self, limit: int = 10, 
                              days: int = 1) -> List[Dict]:
        """Получить топ атакующих IP за указанный период"""
        since_date = datetime.now() - timedelta(days=days)
        
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT ip_address, COUNT(*) as attempts, 
                       MIN(timestamp) as first_attempt,
                       MAX(timestamp) as last_attempt,
                       GROUP_CONCAT(DISTINCT attack_type) as attack_types
                FROM attack_attempts 
                WHERE timestamp > ?
                GROUP BY ip_address
                ORDER BY attempts DESC
                LIMIT ?
            """, (since_date, limit)) as cursor:
                results = await cursor.fetchall()
                
                attackers = []
                for row in results:
                    attackers.append({
                        'ip': row[0],
                        'attempts': row[1],
                        'first_attempt': row[2],
                        'last_attempt': row[3],
                        'attack_types': row[4].split(',') if row[4] else []
                    })
                return attackers
    
    async def get_ip_details(self, ip: str) -> Dict:
        """Получить детальную информацию по IP"""
        async with aiosqlite.connect(self.db_path) as db:
            # Общая статистика по попыткам
            async with db.execute("""
                SELECT COUNT(*) as total_attempts,
                       MIN(timestamp) as first_seen,
                       MAX(timestamp) as last_seen,
                       GROUP_CONCAT(DISTINCT attack_type) as attack_types,
                       GROUP_CONCAT(DISTINCT username) as usernames
                FROM attack_attempts
                WHERE ip_address = ?
            """, (ip,)) as cursor:
                attack_stats = await cursor.fetchone()
            
            # Информация о бане
            async with db.execute("""
                SELECT ban_reason, banned_at, unban_at, is_active, manual_ban
                FROM banned_ips
                WHERE ip_address = ?
                ORDER BY banned_at DESC
                LIMIT 1
            """, (ip,)) as cursor:
                ban_info = await cursor.fetchone()
            
            # Последние попытки
            async with db.execute("""
                SELECT timestamp, attack_type, username, log_line
                FROM attack_attempts
                WHERE ip_address = ?
                ORDER BY timestamp DESC
                LIMIT 10
            """, (ip,)) as cursor:
                recent_attempts = await cursor.fetchall()
            
            return {
                'ip': ip,
                'total_attempts': attack_stats[0] if attack_stats[0] else 0,
                'first_seen': attack_stats[1],
                'last_seen': attack_stats[2],
                'attack_types': attack_stats[3].split(',') if attack_stats[3] else [],
                'usernames': attack_stats[4].split(',') if attack_stats[4] else [],
                'is_banned': ban_info is not None and ban_info[3] == 1,
                'ban_info': {
                    'reason': ban_info[0] if ban_info else None,
                    'banned_at': ban_info[1] if ban_info else None,
                    'unban_at': ban_info[2] if ban_info else None,
                    'manual': bool(ban_info[4]) if ban_info else False
                } if ban_info else None,
                'recent_attempts': [
                    {
                        'timestamp': attempt[0],
                        'type': attempt[1],
                        'username': attempt[2],
                        'log_line': attempt[3]
                    }
                    for attempt in recent_attempts
                ]
            }
    
    async def add_successful_login(self, ip: str, username: str, 
                                 session_info: str = None) -> None:
        """Добавить запись об успешном входе"""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                INSERT INTO successful_logins 
                (ip_address, username, session_info)
                VALUES (?, ?, ?)
            """, (ip, username, session_info))
            await db.commit()
    
    async def get_daily_stats(self) -> Dict:
        """Получить статистику за сегодня"""
        today = datetime.now().date()
        yesterday = today - timedelta(days=1)
        
        async with aiosqlite.connect(self.db_path) as db:
            # Атаки за сегодня
            async with db.execute("""
                SELECT COUNT(*) FROM attack_attempts 
                WHERE DATE(timestamp) = ?
            """, (today,)) as cursor:
                today_attacks = (await cursor.fetchone())[0]
            
            # Уникальные IP за сегодня
            async with db.execute("""
                SELECT COUNT(DISTINCT ip_address) FROM attack_attempts 
                WHERE DATE(timestamp) = ?
            """, (today,)) as cursor:
                today_unique_ips = (await cursor.fetchone())[0]
            
            # Активные баны
            async with db.execute("""
                SELECT COUNT(*) FROM banned_ips 
                WHERE is_active = 1
            """) as cursor:
                active_bans = (await cursor.fetchone())[0]
            
            # Успешные входы за сегодня
            async with db.execute("""
                SELECT COUNT(*) FROM successful_logins 
                WHERE DATE(timestamp) = ?
            """, (today,)) as cursor:
                today_logins = (await cursor.fetchone())[0]
            
            # Сравнение с вчера
            async with db.execute("""
                SELECT COUNT(*) FROM attack_attempts 
                WHERE DATE(timestamp) = ?
            """, (yesterday,)) as cursor:
                yesterday_attacks = (await cursor.fetchone())[0]
            
            return {
                'today': {
                    'attacks': today_attacks,
                    'unique_ips': today_unique_ips,
                    'successful_logins': today_logins
                },
                'yesterday': {
                    'attacks': yesterday_attacks
                },
                'active_bans': active_bans,
                'attack_change': today_attacks - yesterday_attacks
            }
    
    async def cleanup_old_records(self, days: int = 7) -> int:
        """Очистка старых записей"""
        cutoff_date = datetime.now() - timedelta(days=days)
        
        async with aiosqlite.connect(self.db_path) as db:
            # Удаляем старые попытки атак
            async with db.execute("""
                DELETE FROM attack_attempts 
                WHERE timestamp < ?
            """, (cutoff_date,)) as cursor:
                deleted_attempts = cursor.rowcount
            
            # Удаляем неактивные баны старше cutoff_date
            await db.execute("""
                DELETE FROM banned_ips 
                WHERE is_active = 0 AND banned_at < ?
            """, (cutoff_date,))
            
            await db.commit()
            logger.info(f"Очищено {deleted_attempts} старых записей")
            return deleted_attempts
    
    async def is_whitelisted(self, ip: str, whitelist: List[str]) -> bool:
        """Проверка IP в белом списке (поддержка CIDR)"""
        try:
            ip_obj = ipaddress.ip_address(ip)
            for white_ip in whitelist:
                try:
                    # Проверяем как сеть (CIDR)
                    if '/' in white_ip:
                        network = ipaddress.ip_network(white_ip, strict=False)
                        if ip_obj in network:
                            return True
                    # Проверяем как отдельный IP
                    elif ip_obj == ipaddress.ip_address(white_ip):
                        return True
                except Exception:
                    continue
            return False
        except Exception:
            return False

    async def update_daily_stats(self) -> None:
        """Обновить ежедневную статистику"""
        today = datetime.now().date()
        
        async with aiosqlite.connect(self.db_path) as db:
            # Получаем статистику за сегодня
            stats = await self.get_daily_stats()
            
            # Обновляем или создаем запись
            await db.execute("""
                INSERT OR REPLACE INTO daily_stats 
                (date, total_attempts, unique_ips, successful_logins)
                VALUES (?, ?, ?, ?)
            """, (today, stats['today']['attacks'], 
                  stats['today']['unique_ips'], 
                  stats['today']['successful_logins']))
            await db.commit()