#!/usr/bin/env python3
"""Recover the POC's former hostname-keyed JSON documents in place.

Default: read-only report. --apply requires an exclusive backup JSONL path.
Only documents missing @timestamp and containing a valid Telegraf JSON payload
are candidates. Existing IDs are retained; writes use optimistic concurrency.
"""
import argparse
from datetime import datetime, timezone
import json
from pathlib import Path
import urllib.request

BASE = 'http://127.0.0.1:9200'

def request(path, data, method='POST', content_type='application/json'):
    body = data if isinstance(data, bytes) else json.dumps(data).encode()
    req = urllib.request.Request(BASE + path, data=body, method=method, headers={'Content-Type': content_type})
    with urllib.request.urlopen(req, timeout=30) as response:
        return json.load(response)

def recover(source):
    for value in source.values():
        if not isinstance(value, str) or not value.lstrip().startswith('{'):
            continue
        try:
            parsed = json.loads(value)
            if not (isinstance(parsed.get('name'), str) and isinstance(parsed.get('fields'), dict)
                    and isinstance(parsed.get('tags'), dict) and parsed['tags'].get('host')):
                continue
            parsed['@timestamp'] = datetime.fromtimestamp(parsed['timestamp'] / 1000, timezone.utc).isoformat(timespec='milliseconds').replace('+00:00', 'Z')
            if source.get('ingested_at'):
                parsed['ingested_at'] = source['ingested_at']
            return parsed
        except (ValueError, TypeError, KeyError, OverflowError):
            continue
    return None

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--apply', action='store_true')
    parser.add_argument('--backup', type=Path)
    args = parser.parse_args()
    if args.apply and not args.backup:
        parser.error('--apply requires --backup PATH')
    backup = args.backup.open('x') if args.apply else None
    found = changed = skipped = 0
    scroll = None
    try:
        page = request('/host-metrics-v1-*/_search?scroll=2m', {
            'size': 500, 'seq_no_primary_term': True,
            'query': {'bool': {'must_not': [{'exists': {'field': '@timestamp'}}]}}})
        while True:
            scroll = page.get('_scroll_id', scroll)
            hits = page['hits']['hits']
            if not hits:
                break
            operations = []
            for hit in hits:
                parsed = recover(hit['_source'])
                if parsed is None:
                    skipped += 1; continue
                found += 1
                if args.apply:
                    backup.write(json.dumps(hit) + '\n')
                    action = {'index': {'_index': hit['_index'], '_id': hit['_id'],
                                       'if_seq_no': hit['_seq_no'], 'if_primary_term': hit['_primary_term']}}
                    operations.extend([json.dumps(action), json.dumps(parsed)])
            if operations:
                backup.flush()
                import os
                os.fsync(backup.fileno())
                result = request('/_bulk?refresh=wait_for', ('\n'.join(operations) + '\n').encode(), content_type='application/x-ndjson')
                if result.get('errors'):
                    raise RuntimeError('Bulk errors; originals are backed up: ' + json.dumps(result))
                changed += len(result['items'])
            page = request('/_search/scroll', {'scroll': '2m', 'scroll_id': scroll})
    finally:
        if backup:
            backup.close()
        if scroll:
            request('/_search/scroll', {'scroll_id': [scroll]}, method='DELETE')
    print(f'Recoverable: {found}; repaired: {changed}; skipped unrecognized documents: {skipped}')
    if not args.apply:
        print('Read-only report. Use --apply --backup /path/to/original-documents.jsonl to repair.')

if __name__ == '__main__':
    main()