import subprocess
import shlex
import time
import os
import sys
from typing import Dict, Any, List
from src.core.models.hping_test import HpingTestConfig
from src.utils.logger import get_logger

logger = get_logger("hping_tester")

class Hping3Tester:
    def __init__(self, config: HpingTestConfig):
        self.config = config
        self.process = None
        self.start_time = 0
        self.stats = {
            'sent': 0,
            'received': 0,
            'loss_percent': 0.0,
            'duration': 0.0,
            'start_time': 0,
            'end_time': 0
        }
        self.output_file = "/tmp/hping3_output.txt"

    def build_command(self) -> List[str]:
        """Построение команды hping3"""
        cmd = ["hping3"]

        # Тип тестирования
        if self.config.test_type == 'udp':
            cmd.append("--udp")
        elif self.config.test_type == 'icmp':
            cmd.append("--icmp")
        else:  # TCP
            if self.config.flags:
                cmd.append(f"-{self.config.flags}")

        # Порт источника
        if self.config.source_port:
            cmd.extend(["-s", str(self.config.source_port)])

        # Порт назначения
        if self.config.dest_port:
            if self.config.dest_port_range:
                cmd.extend(["-p", f"++{self.config.dest_port}"])
            else:
                cmd.extend(["-p", str(self.config.dest_port)])

        # Спуфинг IP
        if self.config.spoof_ip:
            cmd.extend(["-a", self.config.spoof_ip])

        # Сетевой интерфейс
        if self.config.interface:
            cmd.extend(["--interface", self.config.interface])

        # Сырые данные
        if self.config.raw_data:
            cmd.extend(["--file", self.config.raw_data])
            cmd.extend(["-d", str(self.config.packet_size)])
        else:
            cmd.extend(["-d", str(self.config.packet_size)])

        # Количество пакетов
        if self.config.count:
            cmd.extend(["-c", str(self.config.count)])

        # Режим флуда
        if self.config.flood:
            cmd.append("--flood")
        else:
            cmd.extend(["-i", self.config.interval])

        # Подробный вывод
        if self.config.verbose:
            cmd.append("-V")

        # Целевой адрес
        cmd.append(self.config.target)

        return cmd

    def _is_graphical_environment(self) -> bool:
        """Проверка, доступна ли графическая среда."""
        return bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))

    def _get_terminal_command(self) -> str:
        """Определение команды для терминала kitty."""
        if subprocess.run(["which", "kitty"], capture_output=True).returncode == 0:
            return "kitty"
        logger.error("Terminal 'kitty' not found.")
        return None

    def run(self) -> Dict[str, Any]:
        """Запуск нагрузочного теста в новом терминале kitty"""
        if not self._is_graphical_environment():
            logger.error("Графическая среда недоступна. Запустите тест через терминал с sudo.")
            return self.stats

        command = self.build_command()
        logger.info(f"🚀 Starting hping3 load test: {' '.join(command)}")

        # Подготовка команды для запуска в kitty
        terminal = self._get_terminal_command()
        if not terminal:
            logger.error("❌ Terminal 'kitty' not available. Test aborted.")
            return self.stats

        # Перенаправляем вывод в файл для последующего анализа
        command_str = " ".join([shlex.quote(c) for c in command])
        full_command = f"sudo {command_str} > {self.output_file} 2>&1"

        try:
            self.start_time = time.time()
            self.stats['start_time'] = self.start_time

            # Запускаем команду в kitty
            terminal_cmd = ["kitty", "--", "bash", "-c", full_command]
            logger.debug(f"Executing terminal command: {' '.join(terminal_cmd)}")
            self.process = subprocess.Popen(
                terminal_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )

            # Ожидаем завершения процесса
            while self.process.poll() is None:
                time.sleep(1)  # Проверяем каждую секунду

            self.stats['end_time'] = time.time()
            self.stats['duration'] = self.stats['end_time'] - self.stats['start_time']

            # Читаем вывод из файла
            if os.path.exists(self.output_file):
                with open(self.output_file, 'r') as f:
                    stdout = f.read()
                stderr = ""  # Ошибки также перенаправлены в файл
                os.remove(self.output_file)  # Удаляем временный файл
            else:
                stdout = ""
                stderr = "Output file not found."
                logger.warning("No output file generated by hping3.")

            # Анализ вывода
            self._parse_output(stdout, stderr)
            self._log_summary()

            return self.stats

        except Exception as e:
            logger.error(f"Error running hping3 in kitty: {e}")
            return self.stats
        finally:
            self.process = None

    def stop(self):
        """Остановка нагрузочного теста"""
        if self.process and self.process.poll() is None:
            logger.info("🛑 Stopping hping3 test")
            self.process.terminate()
            try:
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.kill()

    def _parse_output(self, stdout: str, stderr: str):
        """Анализ вывода hping3"""
        if stderr:
            logger.error(f"hping3 errors:\n{stderr}")

        if stdout:
            for line in stdout.split('\n'):
                print(line)
                if "packets tramitted" in line:
                    parts = line.split()
                    try:
                        self.stats['sent'] = int(parts[0])
                        self.stats['received'] = int(parts[3])
                        loss_percent = float(parts[6].replace('%', ''))
                        self.stats['loss_percent'] = loss_percent
                    except (IndexError, ValueError):
                        logger.warning("Failed to parse statistics line")

            logger.debug(f"hping3 output:\n{stdout}")

    def _log_summary(self):
        """Логирование результатов"""
        loss_percent = self.stats.get('loss_percent', 0)
        sent = self.stats.get('sent', 0)
        received = self.stats.get('received', 0)
        duration = self.stats.get('duration', 0)

        if duration > 0:
            pps = sent / duration
        else:
            pps = 0

        logger.success("📊 hping3 test summary:")
        logger.success(f"  Target: {self.config.target}")
        logger.success(f"  Protocol: {self.config.test_type.upper()}")
        logger.success(f"  Packets sent: {sent}")
        logger.success(f"  Packets received: {received}")
        logger.success(f"  Packet loss: {loss_percent}%")
        logger.success(f"  Test duration: {duration:.2f} seconds")
        logger.success(f"  Packets per second: {pps:.2f} pps")
        logger.success(f"  Spoof IP: {self.config.spoof_ip or 'None'}")
        logger.success(f"  Flags: {self.config.flags or 'None'}")
        logger.success(f"  Flood mode: {'Yes' if self.config.flood else 'No'}")
        logger.success(f"  Interval: {self.config.interval if not self.config.flood else 'N/A (flood mode)'}")