#!/usr/bin/env python3
"""Assert end-to-end native host ingestion; bounded wait, nonzero exit on failure."""
import json
import os
import sys
import time
import urllib.error
import urllib.request

BASE = os.environ.get('OPENSEARCH_URL', 'http://127.0.0.1:9200')
REQUIRED = {
    'cpu': 'usage_active', 'mem': 'used_percent', 'disk': 'used_percent',
    'diskio': 'read_bytes_per_sec', 'net': 'bytes_recv_per_sec',
}

def request(path, payload=None):
    r = urllib.request.Request(BASE + path, data=None if payload is None else json.dumps(payload).encode(),
                               headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(r, timeout=20) as response:
        return json.load(response)

failures = []
for attempt in range(36):
    failures = []
    for name, field in REQUIRED.items():
        filters = [{'term': {'name': name}}, {'range': {'@timestamp': {'gte': 'now-2m'}}},
                   {'exists': {'field': 'fields.' + field}}]
        if os.environ.get('METRICS_HOST'):
            filters.append({'term': {'tags.host': os.environ['METRICS_HOST']}})
        if name == 'cpu':
            filters.append({'term': {'tags.cpu': 'cpu-total'}})
        try:
            data = request('/host-metrics-v1-*/_search', {
                'size': 1, 'sort': [{'@timestamp': 'desc'}], 'query': {'bool': {'filter': filters}}})
            hits = data['hits']['hits']
            if not hits:
                failures.append(f'{name}: no fresh {field}'); continue
            doc = hits[0]['_source']
            assert isinstance(doc['fields'][field], (int, float))
            assert all(doc['tags'].get(k) for k in ('host', 'environment', 'role', 'series'))
            assert doc['ingested_at'] and doc['@timestamp']
            if field in ('usage_active', 'used_percent'):
                assert 0 <= doc['fields'][field] <= 100.01
            else:
                assert doc['fields'][field] >= 0
        except (urllib.error.URLError, KeyError, AssertionError) as e:
            failures.append(f'{name}: {e}')
    if not failures:
        break
    if attempt == 0 or attempt % 6 == 0:
        print('Waiting:', '; '.join(failures), flush=True)
    time.sleep(5)
else:
    sys.exit('FAIL: ' + '; '.join(failures) + '\nCheck journalctl -u telegraf-host-metrics and docker compose logs data-prepper.')

caps = request('/host-metrics-v1-*/_field_caps?fields=@timestamp,tags.host,fields.used_percent')
assert 'date' in caps['fields']['@timestamp'], caps
assert 'keyword' in caps['fields']['tags.host'], caps
assert 'double' in caps['fields']['fields.used_percent'], caps
print('PASS: recent native host samples, CPU-total, derived rates, identity tags, timestamps and numeric mappings.')