#!/usr/bin/env python3
"""Send ONE uniquely tagged probe to Kafka, then verify its indexed representation.

Does not reset consumer offsets or write directly to OpenSearch. The measurement
name 'pipeline_probe' is excluded by every supplied metric visualization filter.
"""
import json
from pathlib import Path
import subprocess
import sys
import time
import urllib.error
import urllib.request
import uuid

ROOT = Path(__file__).resolve().parents[1]

def probe():
    probe_id = uuid.uuid4().hex
    sample = {'name': 'pipeline_probe', 'timestamp': int(time.time() * 1000),
              'tags': {'host': 'pipeline-probe', 'environment': 'diagnostic',
                       'role': 'diagnostic', 'series': 'pipeline-probe / diagnostic', 'probe_id': probe_id},
              'fields': {'value': 1.0}}
    print('Publishing one diagnostic measurement to Kafka:', probe_id, flush=True)
    command = ['docker', 'compose', 'exec', '-T', 'kafka', '/opt/kafka/bin/kafka-console-producer.sh',
               '--bootstrap-server', 'kafka:29092', '--topic', 'host-metrics-v1',
               '--producer-property', 'acks=all', '--producer-property', 'delivery.timeout.ms=15000',
               '--producer-property', 'request.timeout.ms=10000', '--producer-property', 'max.block.ms=15000']
    try:
        sent = subprocess.run(command, cwd=ROOT, input=json.dumps(sample) + '\n', text=True,
                              capture_output=True, timeout=40)
    except subprocess.TimeoutExpired:
        print('FAIL: Kafka producer timed out. Check broker/topic/listeners.'); return 2
    if sent.returncode or 'ERROR' in sent.stderr:
        print(sent.stdout, sent.stderr)
        print('FAIL: Kafka did not accept the probe.'); return 2
    print('Producer completed; checking Data Prepper → OpenSearch for up to 60 seconds.', flush=True)
    last_error = ''
    for _ in range(12):
        query = {'size': 1, 'query': {'term': {'tags.probe_id': probe_id}}}
        request = urllib.request.Request('http://127.0.0.1:9200/host-metrics-v1-*/_search',
                                         data=json.dumps(query).encode(), headers={'Content-Type': 'application/json'})
        try:
            with urllib.request.urlopen(request, timeout=5) as response:
                hits = json.load(response)['hits']['hits']
            if hits:
                doc = hits[0]['_source']
                if not doc.get('@timestamp') or doc.get('fields', {}).get('value') != 1:
                    print('FAIL: probe indexed with incorrect transformation:', json.dumps(doc)); return 3
                print('PASS: Kafka → Data Prepper → OpenSearch is working.')
                print('Index:', hits[0]['_index'], 'timestamp:', doc['@timestamp'])
                print('If native host metrics are absent, inspect the Telegraf service/configuration and its Kafka connection.')
                return 0
        except (urllib.error.URLError, KeyError, ValueError) as error:
            last_error = str(error)
        time.sleep(5)
    print('FAIL: probe was not found in OpenSearch within the wait window.', last_error)
    print('Inspect Data Prepper startup errors, consumer lag/backlog, OpenSearch write failures and the DLQ.')
    print('Probe ID (for searching again after a backlog drains):', probe_id)
    return 3

if __name__ == '__main__':
    sys.exit(probe())