#!/usr/bin/env python3
"""
Parallel API Key Rotator Proxy
Rotates Parallel API keys for search/extract/task/monitor endpoints.
Retries on 402 (no credit) and 429 (rate limit) with next key.

Usage:
  python parallel_proxy.py --port 8765 --keys parallel_api_keys.txt
"""

import json, os, sys, time, threading
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse
import requests

PARALLEL_BASE = "https://api.parallel.ai"

class KeyPool:
    def __init__(self, keys_file):
        self.keys = []
        self.lock = threading.Lock()
        self.stats = {}
        self.load_keys(keys_file)
    
    def load_keys(self, keys_file):
        with open(keys_file, 'r') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                key = line.split('#')[0].strip()
                if key:
                    self.keys.append(key)
                    self.stats[key] = {'uses': 0, 'fails': 0, 'last_used': 0}
        print(f"Loaded {len(self.keys)} keys")
    
    def get_key(self):
        with self.lock:
            if not self.keys:
                return None
            key = min(self.keys, key=lambda k: self.stats[k]['uses'])
            self.stats[key]['uses'] += 1
            self.stats[key]['last_used'] = time.time()
            return key
    
    def mark_fail(self, key):
        with self.lock:
            if key in self.stats:
                self.stats[key]['fails'] += 1
    
    def get_stats(self):
        with self.lock:
            total_uses = sum(s['uses'] for s in self.stats.values())
            total_fails = sum(s['fails'] for s in self.stats.values())
            return {
                'total_keys': len(self.keys),
                'total_uses': total_uses,
                'total_fails': total_fails,
            }

class ParallelAPI:
    def __init__(self, key_pool):
        self.pool = key_pool
        self.session = requests.Session()
    
    def _call(self, endpoint, data, method='POST', retry=3):
        for attempt in range(retry):
            key = self.pool.get_key()
            if not key:
                return None, 503, {'error': 'No keys available'}
            
            headers = {
                'Authorization': f'Bearer {key}',
                'Content-Type': 'application/json'
            }
            url = f"{PARALLEL_BASE}{endpoint}"
            
            try:
                if method == 'POST':
                    r = self.session.post(url, headers=headers, json=data, timeout=30)
                else:
                    r = self.session.get(url, headers=headers, params=data, timeout=30)
                
                # 402/429 → try next key
                if r.status_code in (402, 429):
                    self.pool.mark_fail(key)
                    continue
                
                # 200 with credit error body → try next key
                if r.status_code == 200:
                    try:
                        body = r.json()
                        if body.get('type') == 'error' and 'credit' in body.get('error', {}).get('message', '').lower():
                            self.pool.mark_fail(key)
                            continue
                    except:
                        pass
                
                if r.status_code >= 400:
                    self.pool.mark_fail(key)
                    return r.json() if r.text else {}, r.status_code, r.text
                
                return r.json(), r.status_code, None
            except Exception as e:
                self.pool.mark_fail(key)
                continue
        
        return None, 503, {'error': 'All keys exhausted'}
    
    def search(self, queries, **kwargs):
        data = {"search_queries": queries if isinstance(queries, list) else [queries]}
        data.update(kwargs)
        return self._call('/v1/search', data)
    
    def extract(self, urls, **kwargs):
        if isinstance(urls, str):
            urls = [urls]
        return self._call('/v1/extract', {"urls": urls, **kwargs})
    
    def task(self, query, **kwargs):
        return self._call('/v1/task', {"query": query, **kwargs})
    
    def monitor(self, url, **kwargs):
        return self._call('/v1/monitor', {"url": url, **kwargs})