import torch
import torch.nn as nn
import torch.optim as optim
import random
from models.gnn import PhishingGNN_Model
from pipeline.graph_engine import TopologicalGraphEngine
from config import *

def generate_synthetic_warfare(batch_size=100):
    """Generates synthetic network logs simulating both normal traffic and advanced attacks."""
    logs = []
    labels = []
    
    for _ in range(batch_size):
        is_attack = random.random() > 0.5
        
        if is_attack:
            logs.append({
                'ip': f"{random.randint(1, 255)}.{random.randint(1,255)}.0.0",
                'domain': None, 
                'asn': random.choice([666, 9999, 5555]) 
            })
            labels.append([1.0]) 
        else:
            logs.append({
                'ip': f"104.21.{random.randint(1,100)}.{random.randint(1,255)}",
                'domain': f"safe-service-{random.randint(1,100)}.com",
                'asn': 13335 
            })
            labels.append([0.0]) 
            
    return logs, torch.tensor(labels, dtype=torch.float32)

def run_war_games():
    print("[*] Initiating Autonomous Threat War Games...")
    
    in_channels_dict = {'ip': 16, 'domain': 32, 'asn': 8, 'cert': 16}
    model = PhishingGNN_Model(
        metadata=GRAPH_METADATA,
        in_channels_dict=in_channels_dict,
        hidden_channels=HIDDEN_CHANNELS,
        num_heads=NUM_HEADS,
        num_layers=NUM_LAYERS,
        dropout_rate=DROPOUT_RATE
    )
    
    optimizer = optim.AdamW(model.parameters(), lr=0.001)
    criterion = nn.BCELoss()
    engine = TopologicalGraphEngine()

    for wave in range(50):
        logs, labels = generate_synthetic_warfare(batch_size=200)
        x_dict, edge_index_dict = engine.extract_and_build(logs)
        
        optimizer.zero_grad()
        predictions = model(x_dict, edge_index_dict)
        valid_preds = predictions[:len(labels)] 
        
        loss = criterion(valid_preds, labels)
        loss.backward()
        optimizer.step()
        
        if (wave + 1) % 10 == 0:
            print(f"Attack Wave {wave+1:02d}/50 Defeated | Model Penetration Loss: {loss.item():.4f}")

    model.safe_save(MODEL_SAVE_PATH)
    print(f"[+] War games complete. Apex model hardened and saved to: {MODEL_SAVE_PATH}")

if __name__ == "__main__":
    run_war_games()