# src/core/database/db.py
import sqlite3
from pathlib import Path
from typing import Optional
from contextlib import contextmanager

from src.utils.logger import get_logger

logger = get_logger("database")

class Database:
    def __init__(self, db_path: str = "/home/lodqa/attack_module_data/security_scanner.db"):
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.conn: Optional[sqlite3.Connection] = None
        self.init_db()

    def connect(self):
        """Подключение к базе данных."""
        try:
            self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
            self.conn.row_factory = sqlite3.Row
            logger.info(f"Connected to database at {self.db_path}")
        except sqlite3.Error as e:
            logger.error(f"Failed to connect to database: {e}")
            raise

    def init_db(self):
        """Инициализация базы данных и создание таблиц."""
        self.connect()
        cursor = self.conn.cursor()

        # Создание таблиц
        cursor.executescript("""
            CREATE TABLE IF NOT EXISTS sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS scans (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id INTEGER,
                tool TEXT,
                args TEXT,
                summary TEXT,
                duration INTEGER,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS hosts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                scan_id INTEGER,
                ip TEXT,
                mac TEXT,
                os TEXT,
                FOREIGN KEY (scan_id) REFERENCES scans(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS ports (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                host_id INTEGER,
                protocol TEXT,
                port_num INTEGER,
                state TEXT,
                FOREIGN KEY (host_id) REFERENCES hosts(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS service (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                port_id INTEGER,
                name TEXT,
                product TEXT,
                extrainfo TEXT,
                ostype TEXT,
                FOREIGN KEY (port_id) REFERENCES ports(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS cpe (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                service_id INTEGER,
                name TEXT,
                FOREIGN KEY (service_id) REFERENCES service(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS modbus_scan_result (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                scan_id INTEGER,
                active_coils TEXT,
                active_discrete_inputs TEXT,
                active_holding_registers TEXT,
                active_input_registers TEXT,
                FOREIGN KEY (scan_id) REFERENCES scans(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS attacks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id INTEGER,
                tool TEXT,
                args TEXT,
                summary TEXT,
                duration INTEGER,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS attack_result (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                attack_id INTEGER,
                summary TEXT,
                FOREIGN KEY (attack_id) REFERENCES attacks(id) ON DELETE CASCADE
            );
        """)
        self.conn.commit()
        logger.info("Database tables initialized")

    @contextmanager
    def get_cursor(self):
        """Контекстный менеджер для работы с курсором."""
        cursor = self.conn.cursor()
        try:
            yield cursor
            self.conn.commit()
        except Exception as e:
            self.conn.rollback()
            logger.error(f"Database operation failed: {e}")
            raise
        finally:
            cursor.close()

    def close(self):
        """Закрытие соединения с базой данных."""
        if self.conn:
            self.conn.close()
            logger.info("Database connection closed")