from src.core.models.models import Scan, Host, Port, Service, CPE, ModbusScanResult, Attack, AttackResult, Session
from src.core.database.database import Database
from src.utils.logger import get_logger
from typing import List, Optional

logger = get_logger("scanner_service")

class ScannerManager:
    def __init__(self, db: Database):
        self.db = db
    
    def save_session(self, session: Session) -> Optional[Session]:
        """Save a Session object to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO sessions (name, created_at)
                VALUES (?, CURRENT_TIMESTAMP)
                RETURNING *
                """,
                (session.name,)
            )
            row = cursor.fetchone()
            if row:
                return Session(
                    id=row["id"],
                    name=row["name"],
                    created_at=row["created_at"]
                )
        return None

    def get_session_by_id(self, session_id: int) -> Optional[Session]:
        """Retrieve a Session by ID."""
        with self.db.get_cursor() as cursor:
            cursor.execute("SELECT * FROM sessions WHERE id = ?", (session_id,))
            row = cursor.fetchone()
            if row:
                return Session(
                    id=row["id"],
                    name=row["name"],
                    created_at=row["created_at"]
                )
        return None

    def save_scan(self, scan: Scan) -> Optional[Scan]:
        """Save a Scan object to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO scans (session_id, tool, args, summary, duration, created_at)
                VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                RETURNING *
                """,
                (scan.session_id, scan.tool, scan.args, scan.summary, scan.duration)
            )
            row = cursor.fetchone()
            if row:
                return Scan(
                    id=row["id"],
                    session_id=row["session_id"],
                    tool=row["tool"],
                    args=row["args"],
                    summary=row["summary"],
                    duration=row["duration"],
                    created_at=row["created_at"]
                )
        return None
    
    def update_scan(self, scan: Scan) -> bool:
        """Update an existing Scan object."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                UPDATE scans 
                SET summary = ?, duration = ?
                WHERE id = ?
                """,
                (scan.summary, scan.duration, scan.id)
            )
            return cursor.rowcount > 0
    
    def get_scan_by_id(self, scan_id: int) -> Optional[Scan]:
        """Retrieve a Scan by ID."""
        with self.db.get_cursor() as cursor:
            cursor.execute("SELECT * FROM scans WHERE id = ?", (scan_id,))
            row = cursor.fetchone()
            if row:
                return Scan(
                    id=row["id"],
                    session_id=row["session_id"],
                    tool=row["tool"],
                    args=row["args"],
                    summary=row["summary"],
                    duration=row["duration"],
                    created_at=row["created_at"]
                )
        return None
    
    def get_scans_by_session(self, session_id: int) -> List[Scan]:
        """Retrieve all Scans for a session."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                "SELECT * FROM scans WHERE session_id = ? ORDER BY created_at ASC",
                (session_id,)
            )
            rows = cursor.fetchall()
            return [
                Scan(
                    id=row["id"],
                    session_id=row["session_id"],
                    tool=row["tool"],
                    args=row["args"],
                    summary=row["summary"],
                    duration=row["duration"],
                    created_at=row["created_at"]
                ) for row in rows
            ]
    
    def save_host(self, scan_id: int, host: Host) -> Optional[Host]:
        """Save a Host object to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO hosts (scan_id, ip, mac, os)
                VALUES (?, ?, ?, ?)
                RETURNING *
                """,
                (scan_id, host.ip, host.mac, host.os)
            )
            row = cursor.fetchone()
            if row:
                return Host(
                    id=row["id"],
                    scan_id=row["scan_id"],
                    ip=row["ip"],
                    mac=row["mac"],
                    os=row["os"]
                )
        return None
    
    def get_hosts_by_scan(self, scan_id: int) -> List[Host]:
        """Retrieve all Hosts for a scan."""
        with self.db.get_cursor() as cursor:
            cursor.execute("SELECT * FROM hosts WHERE scan_id = ?", (scan_id,))
            rows = cursor.fetchall()
            return [
                Host(
                    id=row["id"],
                    scan_id=row["scan_id"],
                    ip=row["ip"],
                    mac=row["mac"],
                    os=row["os"]
                ) for row in rows
            ]

    def save_port(self, host_id: int, port: Port) -> Optional[Port]:
        """Save a Port object to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO ports (host_id, protocol, port_num, state)
                VALUES (?, ?, ?, ?)
                RETURNING *
                """,
                (host_id, port.protocol, port.port_num, port.state)
            )
            row = cursor.fetchone()
            if row:
                return Port(
                    id=row["id"],
                    host_id=row["host_id"],
                    protocol=row["protocol"],
                    port_num=row["port_num"],
                    state=row["state"]
                )
        return None

    def get_ports_by_host(self, host_id: int) -> List[Port]:
        """Retrieve all Ports for a host."""
        with self.db.get_cursor() as cursor:
            cursor.execute("SELECT * FROM ports WHERE host_id = ?", (host_id,))
            rows = cursor.fetchall()
            return [
                Port(
                    id=row["id"],
                    host_id=row["host_id"],
                    protocol=row["protocol"],
                    port_num=row["port_num"],
                    state=row["state"]
                ) for row in rows
            ]

    def save_service(self, port_id: int, service: Service) -> Optional[Service]:
        """Save a Service object to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO service (port_id, name, product, extrainfo, ostype)
                VALUES (?, ?, ?, ?, ?)
                RETURNING *
                """,
                (port_id, service.name, service.product, service.extrainfo, service.ostype)
            )
            row = cursor.fetchone()
            if row:
                return Service(
                    id=row["id"],
                    port_id=row["port_id"],
                    name=row["name"],
                    product=row["product"],
                    extrainfo=row["extrainfo"],
                    ostype=row["ostype"]
                )
        return None

    def get_services_by_port(self, port_id: int) -> List[Service]:
        """Retrieve all Services for a port."""
        with self.db.get_cursor() as cursor:
            cursor.execute("SELECT * FROM service WHERE port_id = ?", (port_id,))
            rows = cursor.fetchall()
            return [
                Service(
                    id=row["id"],
                    port_id=row["port_id"],
                    name=row["name"],
                    product=row["product"],
                    extrainfo=row["extrainfo"],
                    ostype=row["ostype"]
                ) for row in rows
            ]

    def save_cpe(self, service_id: int, cpe: CPE) -> Optional[CPE]:
        """Save a CPE object to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO cpe (service_id, name)
                VALUES (?, ?)
                RETURNING *
                """,
                (service_id, cpe.name)
            )
            row = cursor.fetchone()
            if row:
                return CPE(
                    id=row["id"],
                    service_id=row["service_id"],
                    name=row["name"]
                )
        return None

    def get_cpes_by_service(self, service_id: int) -> List[CPE]:
        """Retrieve all CPEs for a service."""
        with self.db.get_cursor() as cursor:
            cursor.execute("SELECT * FROM cpe WHERE service_id = ?", (service_id,))
            rows = cursor.fetchall()
            return [
                CPE(
                    id=row["id"],
                    service_id=row["service_id"],
                    name=row["name"]
                ) for row in rows
            ]

    def save_modbus_result(self, scan_id: int, result: dict) -> Optional[ModbusScanResult]:
        """Save Modbus scan results to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO modbus_scan_result (
                    scan_id, active_coils, active_discrete_inputs, 
                    active_holding_registers, active_input_registers
                )
                VALUES (?, ?, ?, ?, ?)
                RETURNING *
                """,
                (
                    scan_id,
                    result.get("active_coils"),
                    result.get("active_discrete_inputs"),
                    result.get("active_holding_registers"),
                    result.get("active_input_registers")
                )
            )
            row = cursor.fetchone()
            if row:
                return ModbusScanResult(
                    id=row["id"],
                    scan_id=row["scan_id"],
                    active_coils=row["active_coils"],
                    active_discrete_inputs=row["active_discrete_inputs"],
                    active_holding_registers=row["active_holding_registers"],
                    active_input_registers=row["active_input_registers"]
                )
        return None
    
    def get_modbus_result(self, scan_id: int) -> Optional[ModbusScanResult]:
        """Retrieve Modbus scan results by scan ID."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                "SELECT * FROM modbus_scan_result WHERE scan_id = ?",
                (scan_id,)
            )
            row = cursor.fetchone()
            if row:
                return ModbusScanResult(
                    id=row["id"],
                    scan_id=row["scan_id"],
                    active_coils=row["active_coils"],
                    active_discrete_inputs=row["active_discrete_inputs"],
                    active_holding_registers=row["active_holding_registers"],
                    active_input_registers=row["active_input_registers"]
                )
        return None

    def save_attack(self, attack: Attack) -> Optional[Attack]:
        """Save an Attack object to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO attacks (session_id, tool, args, summary, duration, created_at)
                VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                RETURNING *
                """,
                (attack.session_id, attack.tool, attack.args, attack.summary, attack.duration)
            )
            row = cursor.fetchone()
            if row:
                return Attack(
                    id=row["id"],
                    session_id=row["session_id"],
                    tool=row["tool"],
                    args=row["args"],
                    summary=row["summary"],
                    duration=row["duration"],
                    created_at=row["created_at"]
                )
        return None

    def update_attack(self, attack: Attack) -> bool:
        """Update an existing Attack object."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                UPDATE attacks 
                SET summary = ?, duration = ?
                WHERE id = ?
                """,
                (attack.summary, attack.duration, attack.id)
            )
            return cursor.rowcount > 0

    def get_attack_by_id(self, attack_id: int) -> Optional[Attack]:
        """Retrieve an Attack by ID."""
        with self.db.get_cursor() as cursor:
            cursor.execute("SELECT * FROM attacks WHERE id = ?", (attack_id,))
            row = cursor.fetchone()
            if row:
                return Attack(
                    id=row["id"],
                    session_id=row["session_id"],
                    tool=row["tool"],
                    args=row["args"],
                    summary=row["summary"],
                    duration=row["duration"],
                    created_at=row["created_at"]
                )
        return None

    def get_attacks_by_session(self, session_id: int) -> List[Attack]:
        """Retrieve all Attacks for a session."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                "SELECT * FROM attacks WHERE session_id = ? ORDER BY created_at DESC",
                (session_id,)
            )
            rows = cursor.fetchall()
            return [
                Attack(
                    id=row["id"],
                    session_id=row["session_id"],
                    tool=row["tool"],
                    args=row["args"],
                    summary=row["summary"],
                    duration=row["duration"],
                    created_at=row["created_at"]
                ) for row in rows
            ]

    def save_attack_result(self, attack_id: int, attack_result: AttackResult) -> Optional[AttackResult]:
        """Save an AttackResult object to the database."""
        with self.db.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO attack_result (attack_id, summary)
                VALUES (?, ?)
                RETURNING *
                """,
                (attack_id, attack_result.summary)
            )
            row = cursor.fetchone()
            if row:
                return AttackResult(
                    id=row["id"],
                    attack_id=row["attack_id"],
                    summary=row["summary"]
                )
        return None

    def get_attack_results_by_attack(self, attack_id: int) -> List[AttackResult]:
        """Retrieve all AttackResults for an attack."""
        with self.db.get_cursor() as cursor:
            cursor.execute("SELECT * FROM attack_result WHERE attack_id = ?", (attack_id,))
            rows = cursor.fetchall()
            return [
                AttackResult(
                    id=row["id"],
                    attack_id=row["attack_id"],
                    summary=row["summary"]
                ) for row in rows
            ]