"""
OmniSec MFT — Managed File Transfer
=====================================
Secure, encrypted, audited file transfer over SFTP/FTPS/SCP.
Supports AES encryption at rest, integrity verification, and
detailed audit logging for compliance workflows.
"""

from __future__ import annotations

import os
import hashlib
import shutil
import time
from pathlib import Path
from typing import Callable, List, Optional, Tuple
from dataclasses import dataclass, field
from omnisec.core.logger import get_logger

log = get_logger("omnisec.mft")


@dataclass
class TransferRecord:
    """Audit record for a single file transfer operation."""
    filename: str
    source: str
    destination: str
    size_bytes: int
    sha256: str
    duration_sec: float
    status: str          # success | failed | skipped
    protocol: str        # sftp | ftps | scp | local
    encrypted: bool
    timestamp: float = field(default_factory=time.time)
    error: Optional[str] = None


class MFTClient:
    """
    Managed File Transfer client with encryption, integrity checks,
    and full audit trail.

    Supports:
      - SFTP (via paramiko)
      - FTPS (via ftplib)
      - Local file operations (copy, move, archive)
      - AES-256 encryption/decryption of files before transfer
      - SHA-256 integrity verification
      - Transfer callbacks for progress reporting
      - Full audit log

    Example:
        >>> client = MFTClient(host="sftp.example.com", username="user", key_path="~/.ssh/id_rsa")
        >>> client.upload("report.pdf", "/remote/reports/report.pdf")
        >>> client.download("/remote/reports/report.pdf", "downloaded.pdf")
    """

    def __init__(
        self,
        host: str = "localhost",
        port: int = 22,
        username: Optional[str] = None,
        password: Optional[str] = None,
        key_path: Optional[str] = None,
        protocol: str = "sftp",
        encrypt_transfers: bool = True,
        encryption_key: Optional[bytes] = None,
    ):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.key_path = key_path
        self.protocol = protocol.lower()
        self.encrypt_transfers = encrypt_transfers
        self._encryption_key = encryption_key
        self._connection = None
        self._audit_log: List[TransferRecord] = []
        log.info(f"MFTClient ready — {protocol.upper()} → {host}:{port}")

    # ── Connection ──────────────────────────────────────────────────

    def connect(self) -> "MFTClient":
        """
        Establish connection to the remote host.

        Returns:
            self (for chaining)

        Raises:
            RuntimeError: If paramiko is not installed for SFTP

        Example:
            >>> client.connect()
        """
        if self.protocol == "sftp":
            self._connect_sftp()
        elif self.protocol == "ftps":
            self._connect_ftps()
        else:
            raise ValueError(f"Unsupported protocol: {self.protocol}. Use 'sftp' or 'ftps'.")
        return self

    def _connect_sftp(self):
        try:
            import paramiko
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            kwargs = {"username": self.username, "port": self.port}
            if self.key_path:
                kwargs["key_filename"] = os.path.expanduser(self.key_path)
            elif self.password:
                kwargs["password"] = self.password
            ssh.connect(self.host, **kwargs)
            self._connection = ssh.open_sftp()
            self._ssh = ssh
            log.info(f"SFTP connected to {self.host}:{self.port}")
        except ImportError:
            raise RuntimeError("paramiko not installed. Run: pip install paramiko")

    def _connect_ftps(self):
        import ftplib
        ftp = ftplib.FTP_TLS()
        ftp.connect(self.host, self.port)
        ftp.login(self.username or "anonymous", self.password or "")
        ftp.prot_p()
        self._connection = ftp
        log.info(f"FTPS connected to {self.host}:{self.port}")

    def disconnect(self) -> None:
        """Close the connection."""
        if self._connection:
            try:
                self._connection.close()
                if hasattr(self, "_ssh"):
                    self._ssh.close()
            except Exception:
                pass
            self._connection = None
            log.info("Connection closed")

    # ── File Hash ───────────────────────────────────────────────────

    @staticmethod
    def sha256(filepath: str, chunk_size: int = 65536) -> str:
        """
        Compute SHA-256 checksum of a local file.

        Args:
            filepath: Path to file
            chunk_size: Read chunk size in bytes

        Returns:
            Hex digest string

        Example:
            >>> MFTClient.sha256("report.pdf")
            'a3f2...'
        """
        h = hashlib.sha256()
        with open(filepath, "rb") as f:
            for chunk in iter(lambda: f.read(chunk_size), b""):
                h.update(chunk)
        return h.hexdigest()

    # ── Encryption ──────────────────────────────────────────────────

    def encrypt_file(self, src: str, dst: Optional[str] = None) -> str:
        """
        AES-256-GCM encrypt a file.

        Args:
            src: Source file path
            dst: Destination path (default: src + '.enc')

        Returns:
            Path to encrypted file

        Raises:
            RuntimeError: If cryptography package not installed

        Example:
            >>> encrypted_path = client.encrypt_file("secret.pdf")
        """
        try:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        except ImportError:
            raise RuntimeError("cryptography not installed. Run: pip install cryptography")

        dst = dst or src + ".enc"
        key = self._encryption_key or os.urandom(32)
        nonce = os.urandom(12)
        aesgcm = AESGCM(key)

        with open(src, "rb") as f:
            data = f.read()
        ciphertext = aesgcm.encrypt(nonce, data, None)

        with open(dst, "wb") as f:
            f.write(nonce + ciphertext)

        log.info(f"Encrypted {src} → {dst}")
        return dst

    def decrypt_file(self, src: str, dst: Optional[str] = None) -> str:
        """
        AES-256-GCM decrypt a file previously encrypted by encrypt_file().

        Args:
            src: Encrypted file path (.enc)
            dst: Output path (default: src without .enc)

        Returns:
            Path to decrypted file

        Example:
            >>> client.decrypt_file("secret.pdf.enc", "secret.pdf")
        """
        try:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        except ImportError:
            raise RuntimeError("cryptography not installed. Run: pip install cryptography")

        dst = dst or (src[:-4] if src.endswith(".enc") else src + ".dec")
        key = self._encryption_key
        if not key:
            raise ValueError("encryption_key must be set for decryption")

        with open(src, "rb") as f:
            payload = f.read()
        nonce, ciphertext = payload[:12], payload[12:]
        aesgcm = AESGCM(key)
        plaintext = aesgcm.decrypt(nonce, ciphertext, None)

        with open(dst, "wb") as f:
            f.write(plaintext)

        log.info(f"Decrypted {src} → {dst}")
        return dst

    # ── Transfer Operations ─────────────────────────────────────────

    def upload(
        self,
        local_path: str,
        remote_path: str,
        callback: Optional[Callable[[int, int], None]] = None,
    ) -> TransferRecord:
        """
        Upload a local file to the remote host.

        Args:
            local_path:  Local file path
            remote_path: Destination path on remote
            callback:    Optional progress callback(bytes_sent, total_bytes)

        Returns:
            TransferRecord with audit details

        Example:
            >>> rec = client.upload("data.csv", "/uploads/data.csv")
            >>> print(rec.sha256)
        """
        start = time.time()
        size = os.path.getsize(local_path)
        checksum = self.sha256(local_path)
        src_to_send = local_path
        encrypted = False

        if self.encrypt_transfers:
            src_to_send = self.encrypt_file(local_path)
            encrypted = True

        try:
            if self.protocol == "sftp":
                self._connection.put(src_to_send, remote_path, callback=callback)
            elif self.protocol == "ftps":
                with open(src_to_send, "rb") as f:
                    self._connection.storbinary(f"STOR {remote_path}", f)

            duration = time.time() - start
            record = TransferRecord(
                filename=os.path.basename(local_path),
                source=local_path, destination=remote_path,
                size_bytes=size, sha256=checksum,
                duration_sec=round(duration, 3),
                status="success", protocol=self.protocol,
                encrypted=encrypted,
            )
            log.info(f"Uploaded {local_path} → {remote_path} ({size} bytes, {duration:.2f}s)")
        except Exception as e:
            record = TransferRecord(
                filename=os.path.basename(local_path),
                source=local_path, destination=remote_path,
                size_bytes=size, sha256=checksum,
                duration_sec=round(time.time() - start, 3),
                status="failed", protocol=self.protocol,
                encrypted=encrypted, error=str(e),
            )
            log.error(f"Upload failed: {e}")
        finally:
            if encrypted and src_to_send != local_path:
                os.unlink(src_to_send)

        self._audit_log.append(record)
        return record

    def download(
        self,
        remote_path: str,
        local_path: str,
        verify_checksum: Optional[str] = None,
    ) -> TransferRecord:
        """
        Download a file from the remote host.

        Args:
            remote_path:     Source path on remote
            local_path:      Local destination path
            verify_checksum: Optional expected SHA-256 to verify integrity

        Returns:
            TransferRecord with audit details

        Example:
            >>> rec = client.download("/uploads/data.csv", "data.csv")
        """
        start = time.time()
        try:
            if self.protocol == "sftp":
                self._connection.get(remote_path, local_path)
            elif self.protocol == "ftps":
                with open(local_path, "wb") as f:
                    self._connection.retrbinary(f"RETR {remote_path}", f.write)

            size = os.path.getsize(local_path)
            checksum = self.sha256(local_path)
            status = "success"
            error = None

            if verify_checksum and checksum != verify_checksum:
                status = "failed"
                error = f"Checksum mismatch: expected {verify_checksum}, got {checksum}"
                log.error(error)

            if self.encrypt_transfers and local_path.endswith(".enc") and self._encryption_key:
                self.decrypt_file(local_path)

            duration = time.time() - start
            record = TransferRecord(
                filename=os.path.basename(local_path),
                source=remote_path, destination=local_path,
                size_bytes=size, sha256=checksum,
                duration_sec=round(duration, 3),
                status=status, protocol=self.protocol,
                encrypted=self.encrypt_transfers, error=error,
            )
            log.info(f"Downloaded {remote_path} → {local_path} ({size} bytes)")
        except Exception as e:
            record = TransferRecord(
                filename=os.path.basename(local_path),
                source=remote_path, destination=local_path,
                size_bytes=0, sha256="",
                duration_sec=round(time.time() - start, 3),
                status="failed", protocol=self.protocol,
                encrypted=self.encrypt_transfers, error=str(e),
            )
            log.error(f"Download failed: {e}")

        self._audit_log.append(record)
        return record

    def list_remote(self, remote_dir: str) -> List[str]:
        """
        List files in a remote directory.

        Args:
            remote_dir: Remote directory path

        Returns:
            List of filenames

        Example:
            >>> files = client.list_remote("/uploads")
        """
        if self.protocol == "sftp":
            return self._connection.listdir(remote_dir)
        elif self.protocol == "ftps":
            return self._connection.nlst(remote_dir)
        return []

    # ── Local Operations ────────────────────────────────────────────

    def local_copy(self, src: str, dst: str) -> TransferRecord:
        """
        Copy a file locally with SHA-256 verification.

        Example:
            >>> client.local_copy("source.csv", "backup/source.csv")
        """
        start = time.time()
        os.makedirs(os.path.dirname(os.path.abspath(dst)), exist_ok=True)
        shutil.copy2(src, dst)
        checksum = self.sha256(dst)
        record = TransferRecord(
            filename=os.path.basename(src),
            source=src, destination=dst,
            size_bytes=os.path.getsize(dst),
            sha256=checksum,
            duration_sec=round(time.time() - start, 4),
            status="success", protocol="local",
            encrypted=False,
        )
        log.info(f"Copied {src} → {dst}")
        self._audit_log.append(record)
        return record

    # ── Audit ───────────────────────────────────────────────────────

    def get_audit_log(self) -> List[TransferRecord]:
        """Return all transfer records."""
        return list(self._audit_log)

    def export_audit_csv(self, path: str) -> None:
        """
        Export audit log to a CSV file.

        Example:
            >>> client.export_audit_csv("audit.csv")
        """
        import csv
        from dataclasses import asdict
        with open(path, "w", newline="") as f:
            if not self._audit_log:
                return
            writer = csv.DictWriter(f, fieldnames=asdict(self._audit_log[0]).keys())
            writer.writeheader()
            for rec in self._audit_log:
                writer.writerow(asdict(rec))
        log.info(f"Audit log exported to {path}")

    def __enter__(self):
        return self.connect()

    def __exit__(self, *_):
        self.disconnect()

    def __repr__(self) -> str:
        return f"MFTClient(protocol={self.protocol!r}, host={self.host!r}, port={self.port})"
