import os
import sys
import json
import time
import urllib.request
import urllib.error

BASE_DIR = '/opt/ai-os/products/ceo/projects/tech/ev_charging_map'
BATCHES_DIR = os.path.join(BASE_DIR, 'batches')
RESULTS_DIR = os.path.join(BASE_DIR, 'results')
MANIFEST_PATH = os.path.join(BASE_DIR, 'batch_manifest.json')
LOG_PATH = os.path.join(BASE_DIR, 'worker.log')

os.makedirs(RESULTS_DIR, exist_ok=True)

def log(msg):
    ts = time.strftime('%Y-%m-%d %H:%M:%S')
    line = f"[{ts}] {msg}"
    print(line)
    with open(LOG_PATH, 'a', encoding='utf-8') as f:
        f.write(line + '\n')

def send_batch(addresses):
    url = 'https://diachi.io/api/convert-batch'
    payload = json.dumps({'addresses': addresses}).encode('utf-8')
    req = urllib.request.Request(url, data=payload, headers={
        'Content-Type': 'application/json',
        'Origin': 'https://diachi.io',
        'Referer': 'https://diachi.io/',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
    })
    
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read().decode('utf-8'))
            return True, data, 0
    except urllib.error.HTTPError as e:
        body = e.read().decode('utf-8', errors='ignore')
        try:
            err_json = json.loads(body)
            retry_after = err_json.get('retryAfter', 180)
            return False, err_json, retry_after
        except Exception:
            return False, {'error': body, 'code': e.code}, 180
    except Exception as e:
        return False, {'error': str(e)}, 60

def run_worker(max_batches=None):
    with open(MANIFEST_PATH, 'r', encoding='utf-8') as f:
        manifest = json.load(f)

    processed_count = 0
    for entry in manifest:
        if max_batches is not None and processed_count >= max_batches:
            log(f"Reached max batches limit ({max_batches}). Stopping worker.")
            break
            
        batch_num = entry['batch_num']
        result_file = os.path.join(RESULTS_DIR, f'batch_{batch_num}_converted.json')
        
        if entry.get('status') == 'completed' and os.path.exists(result_file):
            continue
            
        addr_file = os.path.join(BATCHES_DIR, f'batch_{batch_num}_addresses.json')
        if not os.path.exists(addr_file):
            log(f"Batch file {addr_file} missing, skipping.")
            continue
            
        with open(addr_file, 'r', encoding='utf-8') as f:
            addresses = json.load(f)
            
        log(f"Processing Batch {batch_num} ({len(addresses)} addresses)...")
        
        while True:
            success, resp, retry_after = send_batch(addresses)
            if success:
                log(f"Batch {batch_num} successfully converted!")
                with open(result_file, 'w', encoding='utf-8') as f:
                    json.dump(resp, f, ensure_ascii=False, indent=2)
                
                entry['status'] = 'completed'
                entry['completed_at'] = time.strftime('%Y-%m-%dT%H:%M:%S')
                with open(MANIFEST_PATH, 'w', encoding='utf-8') as f:
                    json.dump(manifest, f, ensure_ascii=False, indent=2)
                    
                processed_count += 1
                # Wait default cooldown of 185s after success before next batch
                log("Waiting cooldown 185s before next batch...")
                time.sleep(185)
                break
            else:
                wait_sec = max(int(retry_after), 10) + 5
                log(f"Rate limited (HTTP 429) or error. Waiting {wait_sec}s before retry. Msg: {resp.get('error')}")
                time.sleep(wait_sec)

if __name__ == '__main__':
    max_b = int(sys.argv[1]) if len(sys.argv) > 1 else None
    run_worker(max_batches=max_b)