#!/usr/bin/env python3
"""
souleyez.storage.exploit_attempts - Track exploitation attempts and outcomes
"""

from datetime import datetime
from typing import Any, Dict, List, Optional

from souleyez.storage.database import get_db

# Valid status values
STATUS_NOT_TRIED = "not_tried"
STATUS_ATTEMPTED = "attempted"
STATUS_FAILED = "failed"
STATUS_SUCCESS = "success"

VALID_STATUSES = [STATUS_NOT_TRIED, STATUS_ATTEMPTED, STATUS_FAILED, STATUS_SUCCESS]


def record_attempt(
    engagement_id: int,
    host_id: int,
    exploit_identifier: str,
    exploit_title: str,
    status: str,
    service_id: Optional[int] = None,
    notes: Optional[str] = None,
) -> int:
    """
    Record an exploit attempt. Creates new record or updates existing one.

    Args:
        engagement_id: ID of the engagement
        host_id: ID of the target host
        exploit_identifier: Unique identifier (e.g., "msf:exploit/unix/ftp/vsftpd_234_backdoor")
        exploit_title: Human-readable exploit title
        status: One of: not_tried, attempted, failed, success
        service_id: Optional service ID this exploit targets
        notes: Optional notes about the attempt

    Returns:
        ID of the created/updated record
    """
    if status not in VALID_STATUSES:
        raise ValueError(f"Invalid status. Must be one of: {', '.join(VALID_STATUSES)}")

    db = get_db()
    conn = db.get_connection()

    # Check if record exists
    existing = conn.execute(
        """
        SELECT id FROM exploit_attempts
        WHERE engagement_id = ? AND host_id = ? AND service_id IS ? AND exploit_identifier = ?
    """,
        (engagement_id, host_id, service_id, exploit_identifier),
    ).fetchone()

    if existing:
        # Update existing record
        attempt_id = existing["id"]
        conn.execute(
            """
            UPDATE exploit_attempts
            SET status = ?,
                exploit_title = ?,
                notes = ?,
                attempted_at = ?,
                updated_at = CURRENT_TIMESTAMP
            WHERE id = ?
        """,
            (status, exploit_title, notes, datetime.now(), attempt_id),
        )
    else:
        # Create new record
        cursor = conn.execute(
            """
            INSERT INTO exploit_attempts (
                engagement_id, host_id, service_id, exploit_identifier,
                exploit_title, status, attempted_at, notes
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        """,
            (
                engagement_id,
                host_id,
                service_id,
                exploit_identifier,
                exploit_title,
                status,
                datetime.now() if status != STATUS_NOT_TRIED else None,
                notes,
            ),
        )
        attempt_id = cursor.lastrowid

    conn.commit()
    conn.close()

    return attempt_id


def get_attempt_status(
    engagement_id: int,
    host_id: int,
    exploit_identifier: str,
    service_id: Optional[int] = None,
) -> Optional[str]:
    """
    Get the status of a specific exploit attempt.

    Returns:
        Status string if attempt exists, None otherwise
    """
    db = get_db()
    conn = db.get_connection()

    result = conn.execute(
        """
        SELECT status FROM exploit_attempts
        WHERE engagement_id = ? AND host_id = ? AND service_id IS ? AND exploit_identifier = ?
    """,
        (engagement_id, host_id, service_id, exploit_identifier),
    ).fetchone()

    conn.close()

    return result["status"] if result else None


def get_attempts_by_host(host_id: int) -> List[Dict[str, Any]]:
    """Get all exploit attempts for a specific host."""
    db = get_db()
    conn = db.get_connection()

    results = conn.execute(
        """
        SELECT
            ea.*,
            s.port,
            s.service_name,
            h.ip_address,
            h.hostname
        FROM exploit_attempts ea
        LEFT JOIN services s ON ea.service_id = s.id
        LEFT JOIN hosts h ON ea.host_id = h.id
        WHERE ea.host_id = ?
        ORDER BY ea.updated_at DESC
    """,
        (host_id,),
    ).fetchall()

    conn.close()

    return [dict(r) for r in results]


def get_attempts_by_engagement(
    engagement_id: int, status_filter: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    Get all exploit attempts for an engagement, optionally filtered by status.

    Args:
        engagement_id: ID of the engagement
        status_filter: Optional status to filter by

    Returns:
        List of attempt records
    """
    db = get_db()
    conn = db.get_connection()

    query = """
        SELECT
            ea.*,
            s.port,
            s.service_name,
            h.ip_address,
            h.hostname
        FROM exploit_attempts ea
        LEFT JOIN services s ON ea.service_id = s.id
        LEFT JOIN hosts h ON ea.host_id = h.id
        WHERE ea.engagement_id = ?
    """
    params = [engagement_id]

    if status_filter:
        if status_filter not in VALID_STATUSES:
            raise ValueError(
                f"Invalid status filter. Must be one of: {', '.join(VALID_STATUSES)}"
            )
        query += " AND ea.status = ?"
        params.append(status_filter)

    query += " ORDER BY ea.updated_at DESC"

    results = conn.execute(query, params).fetchall()

    conn.close()

    return [dict(r) for r in results]


def get_attempt_stats(engagement_id: int) -> Dict[str, Any]:
    """
    Get statistics about exploit attempts for an engagement.

    Returns:
        Dictionary with counts by status and other metrics
    """
    db = get_db()
    conn = db.get_connection()

    # Get status breakdown
    status_counts = conn.execute(
        """
        SELECT status, COUNT(*) as count
        FROM exploit_attempts
        WHERE engagement_id = ?
        GROUP BY status
    """,
        (engagement_id,),
    ).fetchall()

    stats = {"total": 0, "not_tried": 0, "attempted": 0, "failed": 0, "success": 0}

    for row in status_counts:
        stats["total"] += row["count"]
        stats[row["status"]] = row["count"]

    # Get success rate if attempts have been made
    total_attempts = stats["attempted"] + stats["failed"] + stats["success"]
    if total_attempts > 0:
        stats["success_rate"] = (stats["success"] / total_attempts) * 100
    else:
        stats["success_rate"] = 0

    # Get host coverage
    host_coverage = conn.execute(
        """
        SELECT
            COUNT(DISTINCT host_id) as hosts_with_attempts,
            (SELECT COUNT(*) FROM hosts WHERE engagement_id = ?) as total_hosts
        FROM exploit_attempts
        WHERE engagement_id = ?
    """,
        (engagement_id, engagement_id),
    ).fetchone()

    stats["hosts_with_attempts"] = host_coverage["hosts_with_attempts"]
    stats["total_hosts"] = host_coverage["total_hosts"]

    conn.close()

    return stats


def update_attempt_status(
    attempt_id: int, status: str, notes: Optional[str] = None
) -> bool:
    """
    Update the status of an existing attempt.

    Args:
        attempt_id: ID of the attempt record
        status: New status
        notes: Optional notes to add/update

    Returns:
        True if updated, False if not found
    """
    if status not in VALID_STATUSES:
        raise ValueError(f"Invalid status. Must be one of: {', '.join(VALID_STATUSES)}")

    db = get_db()
    conn = db.get_connection()

    cursor = conn.execute(
        """
        UPDATE exploit_attempts
        SET status = ?,
            notes = ?,
            attempted_at = ?,
            updated_at = CURRENT_TIMESTAMP
        WHERE id = ?
    """,
        (
            status,
            notes,
            datetime.now() if status != STATUS_NOT_TRIED else None,
            attempt_id,
        ),
    )

    rows_updated = cursor.rowcount

    conn.commit()
    conn.close()

    return rows_updated > 0


def delete_attempt(attempt_id: int) -> bool:
    """Delete a single attempt record."""
    db = get_db()
    conn = db.get_connection()

    cursor = conn.execute("DELETE FROM exploit_attempts WHERE id = ?", (attempt_id,))
    rows_deleted = cursor.rowcount

    conn.commit()
    conn.close()

    return rows_deleted > 0


def delete_attempts_by_engagement(engagement_id: int) -> int:
    """Delete all attempt records for an engagement. Returns count deleted."""
    db = get_db()
    conn = db.get_connection()

    cursor = conn.execute(
        "DELETE FROM exploit_attempts WHERE engagement_id = ?", (engagement_id,)
    )
    rows_deleted = cursor.rowcount

    conn.commit()
    conn.close()

    return rows_deleted


def delete_attempts_by_host(host_id: int) -> int:
    """Delete all attempt records for a host. Returns count deleted."""
    db = get_db()
    conn = db.get_connection()

    cursor = conn.execute("DELETE FROM exploit_attempts WHERE host_id = ?", (host_id,))
    rows_deleted = cursor.rowcount

    conn.commit()
    conn.close()

    return rows_deleted


def bulk_record_attempts(attempts: List[Dict[str, Any]]) -> int:
    """
    Bulk insert/update multiple attempt records.

    Args:
        attempts: List of dictionaries with keys: engagement_id, host_id, exploit_identifier,
                  exploit_title, status, service_id (optional), notes (optional)

    Returns:
        Number of records created/updated
    """
    if not attempts:
        return 0

    count = 0
    for attempt in attempts:
        record_attempt(
            engagement_id=attempt["engagement_id"],
            host_id=attempt["host_id"],
            exploit_identifier=attempt["exploit_identifier"],
            exploit_title=attempt["exploit_title"],
            status=attempt["status"],
            service_id=attempt.get("service_id"),
            notes=attempt.get("notes"),
        )
        count += 1

    return count


def get_untried_exploits_for_host(engagement_id: int, host_id: int) -> List[str]:
    """
    Get list of exploit identifiers that have not been tried for a host.
    This is a helper for UI filtering - returns only identifiers that either
    have status 'not_tried' or don't have a record at all.

    Returns:
        List of exploit identifiers
    """
    db = get_db()
    conn = db.get_connection()

    results = conn.execute(
        """
        SELECT exploit_identifier
        FROM exploit_attempts
        WHERE engagement_id = ? AND host_id = ? AND status = ?
    """,
        (engagement_id, host_id, STATUS_NOT_TRIED),
    ).fetchall()

    conn.close()

    return [r["exploit_identifier"] for r in results]
