#!/usr/bin/env python3
"""
C2Sentinel Training Script v2 - Improved training with proper normalization
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import random
from tqdm import tqdm
import json

from c2sentinel import (
    C2Sentinel, C2SentinelConfig, LogBERTC2Sentinel,
    FeatureExtractor
)
from safetensors.torch import save_file


class C2TrafficDataset(Dataset):
    """Dataset with normalized features."""

    def __init__(self, num_samples=10000, normalize=True):
        self.samples = []
        self.labels = []
        self.c2_types = []
        self.feature_extractor = FeatureExtractor()

        print(f"Generating {num_samples} training samples...")

        num_c2 = num_samples // 2
        num_benign = num_samples - num_c2

        # Generate C2 samples
        for _ in tqdm(range(num_c2), desc="C2 samples"):
            connections, c2_type = self._generate_c2_traffic()
            features = self.feature_extractor.extract_features(connections)
            self.samples.append(features)
            self.labels.append(1)
            self.c2_types.append(c2_type)

        # Generate benign samples
        for _ in tqdm(range(num_benign), desc="Benign samples"):
            connections = self._generate_benign_traffic()
            features = self.feature_extractor.extract_features(connections)
            self.samples.append(features)
            self.labels.append(0)
            self.c2_types.append(0)

        self.samples = np.array(self.samples, dtype=np.float32)
        self.labels = np.array(self.labels, dtype=np.float32)
        self.c2_types = np.array(self.c2_types, dtype=np.int64)

        # Normalize features (critical for training stability)
        if normalize:
            self.mean = np.mean(self.samples, axis=0)
            self.std = np.std(self.samples, axis=0) + 1e-8
            self.samples = (self.samples - self.mean) / self.std

            # Save normalization params
            np.savez('normalization_params.npz', mean=self.mean, std=self.std)
            print(f"Feature stats - mean range: [{self.mean.min():.2f}, {self.mean.max():.2f}], "
                  f"std range: [{self.std.min():.4f}, {self.std.max():.2f}]")

        # Shuffle
        indices = np.random.permutation(len(self.samples))
        self.samples = self.samples[indices]
        self.labels = self.labels[indices]
        self.c2_types = self.c2_types[indices]

        print(f"C2 samples: {np.sum(self.labels)}, Benign: {len(self.labels) - np.sum(self.labels)}")

    def _generate_c2_traffic(self):
        """Generate C2 beacon traffic with clear patterns."""
        c2_type = random.randint(1, 10)

        # Strong C2 characteristics
        if c2_type <= 3:  # Fast beacon (Metasploit-style)
            interval = random.uniform(2, 15)
            jitter = random.uniform(0, 0.15)  # Low jitter
            port = random.choice([4444, 4445, 5555, 443])
            bytes_sent = random.randint(80, 200)
            bytes_recv = random.randint(40, 150)
        elif c2_type <= 6:  # Medium beacon (Cobalt Strike-style)
            interval = random.uniform(30, 90)
            jitter = random.uniform(0, 0.2)
            port = 443
            bytes_sent = random.randint(60, 150)
            bytes_recv = random.randint(40, 100)
        else:  # Slow beacon (APT-style)
            interval = random.uniform(120, 300)
            jitter = random.uniform(0, 0.1)  # Very low jitter for APT
            port = 443
            bytes_sent = random.randint(50, 120)
            bytes_recv = random.randint(40, 80)

        # Single destination (key C2 indicator)
        dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
        num_connections = random.randint(10, 40)

        connections = []
        timestamp = 1705600000

        for _ in range(num_connections):
            actual_interval = interval * (1 + random.uniform(-jitter, jitter))
            timestamp += actual_interval

            # Very consistent sizes (key C2 indicator)
            size_var = random.uniform(0.95, 1.05)

            connections.append({
                'timestamp': timestamp,
                'dst_ip': dst_ip,
                'dst_port': port,
                'bytes_sent': int(bytes_sent * size_var),
                'bytes_recv': int(bytes_recv * size_var),
                'protocol': 'tcp'
            })

        return connections, c2_type

    def _generate_benign_traffic(self):
        """Generate clearly benign traffic."""
        pattern = random.choice(['browsing', 'api', 'streaming', 'interactive'])

        connections = []
        timestamp = 1705600000

        if pattern == 'browsing':
            # Multiple destinations, highly variable sizes
            for _ in range(random.randint(10, 40)):
                timestamp += random.uniform(0.5, 45)
                connections.append({
                    'timestamp': timestamp,
                    'dst_ip': f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}",
                    'dst_port': random.choice([80, 443]),
                    'bytes_sent': random.randint(200, 5000),
                    'bytes_recv': random.randint(5000, 500000),
                    'protocol': 'tcp'
                })

        elif pattern == 'api':
            # Single dest but HIGHLY variable response sizes
            dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
            for _ in range(random.randint(15, 40)):
                timestamp += random.uniform(0.1, 20)
                connections.append({
                    'timestamp': timestamp,
                    'dst_ip': dst_ip,
                    'dst_port': 443,
                    'bytes_sent': random.randint(100, 3000),
                    'bytes_recv': random.randint(200, 100000),  # Highly variable
                    'protocol': 'tcp'
                })

        elif pattern == 'streaming':
            # Large downloads, irregular timing
            dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
            for _ in range(random.randint(20, 60)):
                timestamp += random.uniform(0.05, 3)
                connections.append({
                    'timestamp': timestamp,
                    'dst_ip': dst_ip,
                    'dst_port': 443,
                    'bytes_sent': random.randint(30, 200),
                    'bytes_recv': random.randint(5000, 150000),
                    'protocol': 'tcp'
                })

        else:  # interactive (ssh-like)
            dst_ip = f"192.168.{random.randint(0,255)}.{random.randint(1,254)}"
            for _ in range(random.randint(15, 50)):
                if random.random() < 0.3:
                    timestamp += random.uniform(3, 45)  # Thinking
                else:
                    timestamp += random.uniform(0.1, 2)  # Typing
                connections.append({
                    'timestamp': timestamp,
                    'dst_ip': dst_ip,
                    'dst_port': 22,
                    'bytes_sent': random.randint(20, 800),
                    'bytes_recv': random.randint(50, 20000),
                    'protocol': 'tcp'
                })

        return connections

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return {
            'features': torch.tensor(self.samples[idx]),
            'label': torch.tensor(self.labels[idx]),
            'c2_type': torch.tensor(self.c2_types[idx])
        }


def train_model(num_epochs=100, batch_size=32, learning_rate=0.0001, num_samples=20000):
    """Train with improved stability."""

    print("=" * 70)
    print("C2Sentinel Model Training v2")
    print("=" * 70)

    config = C2SentinelConfig()
    model = LogBERTC2Sentinel(config)

    # Initialize weights properly
    def init_weights(m):
        if isinstance(m, nn.Linear):
            nn.init.xavier_uniform_(m.weight, gain=0.5)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
    model.apply(init_weights)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    model.to(device)

    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {total_params:,}")

    dataset = C2TrafficDataset(num_samples=num_samples, normalize=True)

    train_size = int(0.9 * len(dataset))
    val_size = len(dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)

    print(f"Train: {train_size}, Val: {val_size}")

    # Simple BCE loss - focus on main task only
    criterion = nn.BCEWithLogitsLoss()

    # Lower LR with warmup
    optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.001)

    # Warmup + cosine decay
    warmup_epochs = 5
    def lr_lambda(epoch):
        if epoch < warmup_epochs:
            return (epoch + 1) / warmup_epochs
        return 0.5 * (1 + np.cos(np.pi * (epoch - warmup_epochs) / (num_epochs - warmup_epochs)))

    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    best_val_acc = 0
    patience = 15
    patience_counter = 0

    for epoch in range(num_epochs):
        model.train()
        train_loss = 0
        train_correct = 0
        train_total = 0

        for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}", leave=False):
            features = batch['features'].to(device)
            labels = batch['label'].to(device)

            optimizer.zero_grad()
            outputs = model(features)

            # Only C2 detection loss
            loss = criterion(outputs['c2_logits'].squeeze(), labels)

            loss.backward()

            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)

            optimizer.step()

            train_loss += loss.item()
            predictions = (torch.sigmoid(outputs['c2_logits'].squeeze()) > 0.5).float()
            train_correct += (predictions == labels).sum().item()
            train_total += labels.size(0)

        scheduler.step()

        # Validation
        model.eval()
        val_correct = 0
        val_total = 0
        val_loss = 0

        with torch.no_grad():
            for batch in val_loader:
                features = batch['features'].to(device)
                labels = batch['label'].to(device)

                outputs = model(features)
                loss = criterion(outputs['c2_logits'].squeeze(), labels)
                val_loss += loss.item()

                predictions = (torch.sigmoid(outputs['c2_logits'].squeeze()) > 0.5).float()
                val_correct += (predictions == labels).sum().item()
                val_total += labels.size(0)

        train_acc = 100 * train_correct / train_total
        val_acc = 100 * val_correct / val_total
        lr = optimizer.param_groups[0]['lr']

        print(f"Epoch {epoch+1}: Loss={train_loss/len(train_loader):.4f}, "
              f"Train={train_acc:.1f}%, Val={val_acc:.1f}%, LR={lr:.6f}")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            patience_counter = 0
            save_file(model.state_dict(), 'c2_sentinel.safetensors')
            print(f"  -> Saved (Val: {val_acc:.1f}%)")
        else:
            patience_counter += 1
            if patience_counter >= patience:
                print(f"Early stopping at epoch {epoch+1}")
                break

    print(f"\nBest validation accuracy: {best_val_acc:.1f}%")
    return model, config


def test_model():
    """Test the trained model."""
    print("\n" + "=" * 70)
    print("Testing Model")
    print("=" * 70)

    sentinel = C2Sentinel.load('c2_sentinel')

    # Test 1: Cobalt Strike
    print("\n[1] Cobalt Strike Beacon (60s interval)...")
    cs = [{'timestamp': 1705600000 + i*60, 'dst_ip': '185.234.72.19', 'dst_port': 443,
           'bytes_sent': 92, 'bytes_recv': 48} for i in range(16)]
    r = sentinel.analyze(cs)
    print(f"  {'✓ C2 DETECTED' if r.is_c2 else '✗ No C2'} (prob={r.c2_probability:.2%})")

    # Test 2: Metasploit
    print("\n[2] Metasploit Beacon (5s interval, port 4444)...")
    msf = [{'timestamp': 1705600000 + i*5, 'dst_ip': '10.10.10.10', 'dst_port': 4444,
            'bytes_sent': 150, 'bytes_recv': 400} for i in range(20)]
    r = sentinel.analyze(msf)
    print(f"  {'✓ C2 DETECTED' if r.is_c2 else '✗ No C2'} (prob={r.c2_probability:.2%})")

    # Test 3: Slow APT beacon
    print("\n[3] APT Slow Beacon (120s interval)...")
    apt = [{'timestamp': 1705600000 + i*120, 'dst_ip': '45.33.32.156', 'dst_port': 443,
            'bytes_sent': 80, 'bytes_recv': 60} for i in range(12)]
    r = sentinel.analyze(apt)
    print(f"  {'✓ C2 DETECTED' if r.is_c2 else '✗ No C2'} (prob={r.c2_probability:.2%})")

    # Test 4: Web browsing (should be benign)
    print("\n[4] Web Browsing (should be clean)...")
    browse = [{'timestamp': 1705600000 + i*random.uniform(2, 30),
               'dst_ip': f"{random.randint(1,200)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}",
               'dst_port': 443, 'bytes_sent': random.randint(500, 3000),
               'bytes_recv': random.randint(10000, 500000)} for i in range(20)]
    r = sentinel.analyze(browse)
    print(f"  {'✗ C2 DETECTED (FP!)' if r.is_c2 else '✓ Clean'} (prob={r.c2_probability:.2%})")

    # Test 5: SSH keepalive
    print("\n[5] SSH Keepalive (should be clean)...")
    ssh = [{'timestamp': 1705600000 + i*30, 'dst_ip': '192.168.1.50', 'dst_port': 22,
            'bytes_sent': 48, 'bytes_recv': 48} for i in range(15)]
    r = sentinel.analyze(ssh)
    print(f"  {'✗ C2 DETECTED (FP!)' if r.is_c2 else '✓ Clean'} (prob={r.c2_probability:.2%})")
    print(f"  Pattern: {r.matched_legitimate_pattern}")

    # Test 6: API calls (should be benign)
    print("\n[6] API Calls (should be clean)...")
    api = [{'timestamp': 1705600000 + i*random.uniform(0.5, 10),
            'dst_ip': '52.85.132.99', 'dst_port': 443,
            'bytes_sent': random.randint(100, 2000),
            'bytes_recv': random.randint(500, 80000)} for i in range(25)]
    r = sentinel.analyze(api)
    print(f"  {'✗ C2 DETECTED (FP!)' if r.is_c2 else '✓ Clean'} (prob={r.c2_probability:.2%})")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--samples', type=int, default=20000)
    parser.add_argument('--batch-size', type=int, default=32)
    parser.add_argument('--lr', type=float, default=0.0001)
    parser.add_argument('--test-only', action='store_true')
    args = parser.parse_args()

    if args.test_only:
        test_model()
    else:
        train_model(num_epochs=args.epochs, batch_size=args.batch_size,
                   learning_rate=args.lr, num_samples=args.samples)
        test_model()