"""Gradio demo for DGA domain classifier on HuggingFace Spaces.

This demo loads the model from HuggingFace Hub and provides an interactive
interface for classifying domains as legitimate or DGA-generated.
"""

import torch
import gradio as gr

# Import custom model and encoding
from model import DGAEncoderForSequenceClassification
from charset import encode_domain


# Load model from HuggingFace Hub
MODEL_NAME = "ccss17/dga-transformer-encoder"
print(f"Loading model from {MODEL_NAME}...")
model = DGAEncoderForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
print(f"Model loaded on {device}")


def predict_domain(domain: str):
    """Classify a domain as legitimate or DGA-generated.

    Args:
        domain: Domain name to classify (e.g., "google.com", "xjkd8f2h.com")

    Returns:
        tuple: (prediction_label, confidence_score, html_output)
    """
    if not domain or not domain.strip():
        return "⚠️ Invalid Input", "", "Please enter a domain name."

    domain = domain.strip().lower()

    # Encode domain to token IDs
    input_ids = torch.tensor(
        [encode_domain(domain, max_len=64)], device=device
    )

    # Get model prediction
    with torch.no_grad():
        outputs = model(input_ids=input_ids)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=-1)
        pred_class = torch.argmax(probs, dim=-1).item()
        confidence = probs[0, pred_class].item()

    # Format output
    label_names = ["✅ Legitimate", "🚨 DGA (Malicious)"]
    prediction = label_names[pred_class]
    confidence_pct = f"{confidence * 100:.2f}%"

    # Create detailed HTML output
    legit_prob = probs[0, 0].item()
    dga_prob = probs[0, 1].item()

    html_output = f"""
    <div style="padding: 20px; border-radius: 10px; background: {"#d4edda" if pred_class == 0 else "#f8d7da"};">
        <h2 style="margin-top: 0; color: {"#155724" if pred_class == 0 else "#721c24"};">
            {prediction}
        </h2>
        <p style="font-size: 18px; margin: 10px 0;">
            <strong>Domain:</strong> <code>{domain}</code>
        </p>
        <p style="font-size: 18px; margin: 10px 0;">
            <strong>Confidence:</strong> {confidence_pct}
        </p>
        <hr style="margin: 15px 0; border: none; border-top: 1px solid #ccc;">
        <h3>Probability Breakdown:</h3>
        <div style="margin: 10px 0;">
            <div style="display: flex; justify-content: space-between; margin-bottom: 5px;">
                <span>✅ Legitimate:</span>
                <strong>{legit_prob * 100:.2f}%</strong>
            </div>
            <div style="background: #e9ecef; border-radius: 5px; overflow: hidden; height: 20px;">
                <div style="background: #28a745; height: 100%; width: {legit_prob * 100}%;"></div>
            </div>
        </div>
        <div style="margin: 10px 0;">
            <div style="display: flex; justify-content: space-between; margin-bottom: 5px;">
                <span>🚨 DGA (Malicious):</span>
                <strong>{dga_prob * 100:.2f}%</strong>
            </div>
            <div style="background: #e9ecef; border-radius: 5px; overflow: hidden; height: 20px;">
                <div style="background: #dc3545; height: 100%; width: {dga_prob * 100}%;"></div>
            </div>
        </div>
    </div>
    """

    return prediction, confidence_pct, html_output


def predict_batch(domains_text: str):
    """Classify multiple domains at once.

    Args:
        domains_text: Newline-separated list of domains

    Returns:
        str: Formatted results table in Markdown
    """
    if not domains_text or not domains_text.strip():
        return "Please enter one or more domain names (one per line)."

    domains = [
        d.strip() for d in domains_text.strip().split("\n") if d.strip()
    ]

    if not domains:
        return "No valid domains provided."

    results = []
    results.append("| Domain | Prediction | Confidence |")
    results.append("|--------|------------|------------|")

    for domain in domains:
        domain = domain.lower()

        # Encode domain
        input_ids = torch.tensor(
            [encode_domain(domain, max_len=64)], device=device
        )

        # Predict
        with torch.no_grad():
            outputs = model(input_ids=input_ids)
            logits = outputs.logits
            probs = torch.softmax(logits, dim=-1)
            pred_class = torch.argmax(probs, dim=-1).item()
            confidence = probs[0, pred_class].item()

        label = "Legitimate ✅" if pred_class == 0 else "DGA 🚨"
        confidence_pct = f"{confidence * 100:.1f}%"

        results.append(f"| `{domain}` | {label} | {confidence_pct} |")

    return "\n".join(results)


# Example domains for quick testing
EXAMPLES = [
    ["google.com"],
    ["github.com"],
    ["stackoverflow.com"],
    ["xjkd8f2h.com"],
    ["qwfp93nx.net"],
    ["h4fk29fd.org"],
    ["facebook.com"],
    ["fjdkslajf.com"],
]

BATCH_EXAMPLES = [
    "google.com\ngithub.com\nstackoverflow.com",
    "xjkd8f2h.com\nqwfp93nx.net\nh4fk29fd.org",
    "amazon.com\nfacebook.com\ntwitter.com\nmicrosoft.com",
]


# Create Gradio interface
with gr.Blocks(title="DGA Domain Classifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🔍 DGA Domain Classifier
    
    **Detect malicious domains generated by Domain Generation Algorithms (DGAs)**
    
    This model uses a transformer-based neural network to classify domains as either:
    - ✅ **Legitimate**: Normal, human-registered domains
    - 🚨 **DGA (Malicious)**: Algorithmically-generated domains used by malware for C2 communication
    
    ---
    
    ### Model Details
    - **Architecture**: Custom Transformer Encoder (4 layers, 256 dim)
    - **Parameters**: 3.2M
    - **Accuracy**: 96.78% F1 score on test set
    - **Inference Speed**: <1ms per domain
    
    ---
    """)

    with gr.Tabs():
        # Tab 1: Single domain prediction
        with gr.Tab("Single Domain"):
            gr.Markdown("### Test a single domain name")

            with gr.Row():
                with gr.Column(scale=2):
                    domain_input = gr.Textbox(
                        label="Enter Domain Name",
                        placeholder="e.g., google.com, xjkd8f2h.com",
                        lines=1,
                    )
                    predict_btn = gr.Button(
                        "🔍 Classify Domain", variant="primary", size="lg"
                    )

                with gr.Column(scale=1):
                    prediction_output = gr.Textbox(label="Prediction", lines=1)
                    confidence_output = gr.Textbox(label="Confidence", lines=1)

            html_output = gr.HTML(label="Detailed Results")

            predict_btn.click(
                fn=predict_domain,
                inputs=domain_input,
                outputs=[prediction_output, confidence_output, html_output],
            )

            gr.Markdown("### Try these examples:")
            gr.Examples(
                examples=EXAMPLES,
                inputs=domain_input,
                outputs=[prediction_output, confidence_output, html_output],
                fn=predict_domain,
                cache_examples=False,
            )

        # Tab 2: Batch prediction
        with gr.Tab("Batch Prediction"):
            gr.Markdown("### Classify multiple domains at once (one per line)")

            batch_input = gr.Textbox(
                label="Enter Domains (one per line)",
                placeholder="google.com\ngithub.com\nxjkd8f2h.com",
                lines=8,
            )
            batch_btn = gr.Button(
                "🔍 Classify All Domains", variant="primary", size="lg"
            )
            batch_output = gr.Markdown(label="Results")

            batch_btn.click(
                fn=predict_batch,
                inputs=batch_input,
                outputs=batch_output,
            )

            gr.Markdown("### Try these examples:")
            gr.Examples(
                examples=BATCH_EXAMPLES,
                inputs=batch_input,
                outputs=batch_output,
                fn=predict_batch,
                cache_examples=False,
            )

        # Tab 3: About
        with gr.Tab("About"):
            gr.Markdown("""
            ## 📚 What are DGAs?
            
            **Domain Generation Algorithms (DGAs)** are techniques used by malware to generate 
            large numbers of pseudo-random domain names for C2 (command-and-control) communication.
            
            ### Why DGAs are dangerous:
            - **Evasion**: Traditional blacklists can't keep up with thousands of generated domains
            - **Resilience**: Even if some domains are blocked, malware can try others
            - **Stealth**: DGA domains look random, making detection challenging
            
            ### How this model works:
            
            1. **Character-level tokenization**: Breaks domain into individual characters
            2. **Transformer encoder**: Learns patterns in character sequences
            3. **Self-attention**: Detects unusual character combinations (e.g., `xqz`, `fgh`)
            4. **Classification**: Predicts if domain is legitimate or DGA-generated
            
            ### Key Features:
            - **High accuracy**: 96.78% F1 score on test set
            - **Fast inference**: <1ms per domain (GPU) or ~10ms (CPU)
            - **Lightweight**: Only 3.2M parameters
            - **Production-ready**: Trained on real-world malware domains
            
            ### Examples:
            
            **Legitimate domains** (structured, pronounceable):
            - `google.com`, `github.com`, `stackoverflow.com`
            - `api-docs.company.com`, `cdn-assets.example.org`
            
            **DGA domains** (random, unpronounceable):
            - `xjkd8f2h.com`, `qwfp93nx.net`, `h4fk29fd.org`
            - `kdjf92jd.info`, `zmxbv73k.biz`
            
            ---
            
            ### Technical Details:
            - **Model**: Custom Transformer Encoder
            - **Training data**: ExtraHop DGA dataset
            - **Framework**: PyTorch + HuggingFace Transformers
            - **Experiment tracking**: Weights & Biases
            
            ### Links:
            - [Model Card](https://huggingface.co/ccss17/dga-transformer-encoder)
            - [GitHub Repository](https://github.com/ccss17/DGA-Transformer-Encoder)
            - [ExtraHop DGA Dataset](https://github.com/extrahop/dga-training-data)
            
            ---
            
            **Built with ❤️ using PyTorch, HuggingFace, and Gradio**
            """)

    gr.Markdown("""
    ---
    <div style="text-align: center; color: #666; font-size: 14px;">
        <p>⚠️ <strong>Disclaimer</strong>: This model is for educational and research purposes. 
        Always use multiple detection methods in production security systems.</p>
        <p>Model accuracy: 96.78% | False positive rate: ~3%</p>
    </div>
    """)


if __name__ == "__main__":
    demo.launch()