import json
import uuid
import random
import time
from collections import Counter
from kafka import KafkaProducer


KAFKA_TOPIC = "ivt_requests"
KAFKA_SERVER = "localhost:9092"

producer = KafkaProducer(
    bootstrap_servers=KAFKA_SERVER,
    value_serializer=lambda v: json.dumps(v).encode("utf-8")
)


BLOCKLISTED_IPS = [
    "192.0.2.55",
    "10.10.10.10",
    "203.0.113.25"
]

BAD_UA = "python-requests/2.31.0"

CLEAN_UAS = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/126.0.0.0 Safari/537.36",

    "Mozilla/5.0 (Macintosh; Intel Mac OS X 14_5) "
    "AppleWebKit/605.1.15 (KHTML, like Gecko) "
    "Version/17.5 Safari/605.1.15",

    "Mozilla/5.0 (iPhone; CPU iPhone OS 17_5 like Mac OS X) "
    "AppleWebKit/605.1.15 (KHTML, like Gecko) "
    "Version/17.5 Mobile/15E148 Safari/604.1"
]


def random_ip():
    return (
        f"{random.randint(20, 223)}."
        f"{random.randint(0, 255)}."
        f"{random.randint(0, 255)}."
        f"{random.randint(1, 254)}"
    )


SUSPICIOUS_IP_POOL = [
    f"198.51.100.{n}"
    for n in range(60, 70)
]


def create_openrtb_request(ip, user_agent, tag):

    return {
        "id": f"{tag}_{uuid.uuid4().hex[:8]}",

        "imp": [
            {
                "id": "1",
                "banner": {
                    "format": [
                        {
                            "w": 300,
                            "h": 250
                        }
                    ]
                }
            }
        ],

        "device": {
            "ip": ip,
            "ua": user_agent,
            "geo": {
                "country": "US",
                "region": "CA",
                "city": "Los Angeles"
            }
        },

        "user": {
            "id": f"user_{uuid.uuid4().hex[:8]}"
        },

        "site": {
            "domain": "example.com"
        },

        "at": 1
    }


def send_request(request_data):

    producer.send(
        KAFKA_TOPIC,
        request_data
    )


def warm_up_state():

    print("Warming up blocklisted IPs...")

    for ip in BLOCKLISTED_IPS:

        for _ in range(3):

            send_request(
                create_openrtb_request(
                    ip,
                    BAD_UA,
                    "warmup"
                )
            )

            time.sleep(0.15)

    print("Warming up suspicious IP pool...")

    for ip in SUSPICIOUS_IP_POOL:

        for _ in range(2):

            send_request(
                create_openrtb_request(
                    ip,
                    BAD_UA,
                    "warmup"
                )
            )

            time.sleep(0.15)

    producer.flush()


def build_batch(
    total,
    clean_pct,
    suspicious_pct,
    blocked_pct
):

    n_clean = round(
        total * clean_pct / 100
    )

    n_suspicious = round(
        total * suspicious_pct / 100
    )

    n_blocked = (
        total
        - n_clean
        - n_suspicious
    )

    plan = (
        ["clean"] * n_clean
        + ["suspicious"] * n_suspicious
        + ["blocked"] * n_blocked
    )

    random.shuffle(plan)

    return (
        plan,
        n_clean,
        n_suspicious,
        n_blocked
    )


def make_request_for(category):

    if category == "clean":

        ip = random_ip()
        ua = random.choice(CLEAN_UAS)

    elif category == "suspicious":

        ip = random.choice(
            SUSPICIOUS_IP_POOL
        )

        ua = BAD_UA

    else:

        ip = random.choice(
            BLOCKLISTED_IPS
        )

        ua = BAD_UA

    return create_openrtb_request(
        ip,
        ua,
        category
    )


def run_batch(
    total=100,
    clean_pct=50,
    suspicious_pct=30,
    blocked_pct=20
):

    assert (
        clean_pct
        + suspicious_pct
        + blocked_pct
        == 100
    )

    warm_up_state()

    plan, n_clean, n_suspicious, n_blocked = build_batch(
        total,
        clean_pct,
        suspicious_pct,
        blocked_pct
    )

    print(
        f"\nGenerating {total} requests -> "
        f"{n_clean} clean, "
        f"{n_suspicious} suspicious, "
        f"{n_blocked} blocked"
    )

    print("=" * 70)

    for i, category in enumerate(plan, 1):

        request = make_request_for(category)

        send_request(request)

        print(
            f"[{i:>3}/{total}] "
            f"sent={category:<10} "
            f"id={request['id']}"
        )

        time.sleep(0.05)

    producer.flush()

    print("\n" + "=" * 70)
    print("All requests sent to Kafka")
    print(f"Topic: {KAFKA_TOPIC}")
    print(f"Total requests: {total}")
    print(
        f"Clean: {n_clean} | "
        f"Suspicious: {n_suspicious} | "
        f"Blocked: {n_blocked}"
    )

    producer.close()


if __name__ == "__main__":

    run_batch(
        total=100,
        clean_pct=50,
        suspicious_pct=30,
        blocked_pct=20
    )