import os
import random
import time
import threading
import requests
import json
from flask import Flask, request, jsonify, send_file, make_response
from flask_cors import CORS

app = Flask(__name__)
# Enable CORS secara explicit dan permissive
CORS(app, resources={r"/*": {"origins": "*"}})

# --- TARGET CONFIG ---
# URL Space Flux kamu
BASE_URL = "https://black-forest-labs-flux-2-dev.hf.space"

# --- TOKEN LIST (HARDCODED) ---
RAW_TOKENS = [
    "hf_+PiRCDDtPcPFMLWkTkVaZmzoleHOunXnLIA",
    "hf_+BHvZXGICstaktSwycmwNmzHGrTNmKxnlRZ",
    "hf_+ZdgawyTPzXIpwhnRYIteUKSMsWnEDtGKtM",
    "hf_+nMiFYAFsINxAJWPwiCQlaunmdgmrcxKoaT",
    "hf_+PccpUIbTckCiafwErDLkRlsvqhgtfZaBHL",
    "hf_+faGyXBPfBkaHXDMUSJtxEggonhhZbomFIz",
    "hf_+SndsPaRWsevDXCgZcSjTUlBYUJqOkSfFmn",
    "hf_+CqobFdUpeVCeuhUaiuXwvdczBUmoUHXRGa",
    "hf_+JKCQYUhhHPPkpucegqkNSyureLdXpmeXRF",
    "hf_+tBYfslUwHNiNMufzwAYIlrDVovEWmOQulC",
    "hf_+LKLdrdUxyUyKODSUthmqHXqDMfHrQueera",
    "hf_+ivSBboJYQVcifWkCNcOTOnxUQrZOtOglnU"
]

class TokenManager:
    def __init__(self, raw_tokens):
        self.tokens = [t.replace("+", "").strip() for t in raw_tokens if t.strip()]
        self.current_index = 0
        self.lock = threading.Lock()

    def get_token(self):
        with self.lock:
            if not self.tokens: raise Exception("No tokens available")
            token = self.tokens[self.current_index]
            self.current_index = (self.current_index + 1) % len(self.tokens)
            return token

token_manager = TokenManager(RAW_TOKENS)

def request_flux_gradio4(prompt, width, height, guidance_scale, seed):
    """
    Menangani request khusus untuk Gradio 4 (Flux Standard)
    Endpoint: /gradio_api/call/infer
    """
    max_retries = 5
    attempt = 0
    last_error = ""
    
    # Endpoint Gradio 4 yang Benar
    api_url = f"{BASE_URL}/gradio_api/call/infer"

    while attempt < max_retries:
        token = token_manager.get_token()
        print(f"[LOG] Attempt {attempt+1} - Token: ...{token[-5:]}")
        
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }

        # Payload Gradio 4 (Tanpa fn_index, langsung data array)
        payload = {
            "data": [
                prompt,
                int(seed) if seed is not None else 0,
                True if seed is None else False, # randomize_seed
                int(width),
                int(height),
                float(guidance_scale),
                28 # num_inference_steps
            ]
        }

        try:
            # 1. Kirim Job ID Request
            response = requests.post(api_url, headers=headers, json=payload, timeout=30)
            
            if response.status_code != 200:
                raise Exception(f"HTTP Error {response.status_code}: {response.text}")

            resp_json = response.json()
            # Gradio 4 mengembalikan { "event_id": "..." }
            event_id = resp_json.get("event_id")
            if not event_id:
                raise Exception(f"No event_id returned: {resp_json}")

            # 2. Polling Hasil (Gradio 4 butuh polling status request)
            # Kita tembak endpoint result untuk event_id tersebut
            result_url = f"{BASE_URL}/gradio_api/call/infer/{event_id}"
            
            # Tunggu proses selesai (polling loop)
            # Karena flux cepat, kita poll tiap 1 detik
            for _ in range(60): # Max tunggu 60 detik
                time.sleep(1) 
                poll_resp = requests.get(result_url, headers=headers, stream=True) # Stream true untuk jaga2
                
                # Gradio 4 streaming response berbentuk text event stream
                # Kita cari baris yang mengandung "complete" atau data url
                content = poll_resp.text
                
                if "process_completed" in content and "url" in content:
                    # Parsing manual dari format event-stream
                    # Biasanya format: event: complete \n data: [...]
                    lines = content.split('\n')
                    for line in lines:
                        if line.startswith('data:'):
                            try:
                                data_block = json.loads(line[5:])
                                # Cari URL output
                                if isinstance(data_block, list) and len(data_block) > 0:
                                    file_info = data_block[0]
                                    final_url = file_info.get("url")
                                    if final_url:
                                        # Download Gambar Akhir
                                        print(f"[LOG] Downloading: {final_url}")
                                        img_res = requests.get(final_url, timeout=60)
                                        temp_filename = f"/tmp/flux_{int(time.time())}_{random.randint(0,999)}.webp"
                                        with open(temp_filename, 'wb') as f:
                                            f.write(img_res.content)
                                        return temp_filename
                            except:
                                continue
            
            raise Exception("Timeout waiting for Gradio 4 result")

        except Exception as e:
            err_msg = str(e)
            print(f"[ERROR] {err_msg}")
            if "429" in err_msg or "quota" in err_msg.lower():
                print("Rate limit detected.")
            else:
                last_error = err_msg
            attempt += 1
            time.sleep(1)

    raise Exception(f"Gagal Total: {last_error}")

# --- ROUTES ---

@app.route('/', methods=['GET'])
def home():
    return jsonify({"status": "Online", "mode": "Gradio 4 Native"})

# Handle OPTIONS request manual untuk mencegah Failed to Fetch
@app.route('/generate', methods=['OPTIONS'])
def options_generate():
    response = make_response()
    response.headers.add("Access-Control-Allow-Origin", "*")
    response.headers.add("Access-Control-Allow-Headers", "*")
    response.headers.add("Access-Control-Allow-Methods", "POST, OPTIONS")
    return response

@app.route('/generate', methods=['POST'])
def generate():
    data = request.json
    if not data or 'prompt' not in data:
        return jsonify({"error": "Prompt required"}), 400

    prompt = data.get('prompt')
    width = data.get('width', 1024)
    height = data.get('height', 1024)
    guidance = data.get('guidance', 3.5)
    seed = data.get('seed', None)

    try:
        image_path = request_flux_gradio4(prompt, width, height, guidance, seed)
        return send_file(image_path, mimetype='image/webp')
    except Exception as e:
        # Return JSON error tapi dengan CORS header agar terbaca di frontend
        resp = jsonify({"error": str(e)})
        resp.headers.add("Access-Control-Allow-Origin", "*")
        return resp, 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)