import optuna
import requests
import concurrent.futures
import queue
import os
import glob
import numpy as np
import hashlib
import time
import json
import threading
import math
from google import genai
from google.genai import types

# ==============================================================================
# CONFIGURATION
# ==============================================================================
API_TOKENS = [
    ""
]
GEMINI_API_KEY = ""
INITIAL_PRESET = {
    "temp": 1.0681743378736395, 
    "top_p": 0.8517205306321967, 
    "min_p": 0.006728698878705197, 
    "freq": 0.08057139543156233, 
    "pres": 0.45539559524058393, 
    "top_k": 42317
}
MAX_SAMPLES_PER_TRIAL = 305
EMBEDDING_MODEL = "gemini-embedding-2-preview"
NAI_MODEL = "xialong-v1"
LOGIT_BIAS = {
  "151331": -100,
  "151336": -100,
  "151350": -100,
  "151351": -100,
  "151360": -100
}
PROMPT_TOKENS = [
    151331, 151333, 151335, 198, 2610, 525, 1599, 530, 644, 320, 99650, 82225, 235, 701, 458,
    15223, 1614, 1875, 295, 48576, 553, 1527, 5524, 276, 13, 1446, 1795, 279, 1196, 594, 11216,
    23573, 1393, 12672, 27505, 11, 10903, 681, 11, 323, 7989, 311, 1449, 2033, 13, 58030, 697,
    7743, 323, 1707, 311, 2432, 1128, 279, 3383, 18126, 13, 151336, 198, 7984, 13, 151360, 151337,
    198, 151350, 151351, 198, 30220
] # 30220 = "What
CONTEXT_PREFIX = "\"What" # (1 token, 1024+1 = 1025 token ai content / 1025 token human content)

# ==============================================================================
# MISCELLANEOUS GLOBAL VARIABLES
# ==============================================================================

API_URL = "https://text.novelai.net/oa/v1/completions"

try:
    gemini_client = genai.Client(api_key=GEMINI_API_KEY)
except Exception as e:
    print(f"Failed to initialize Gemini Client. Make sure GEMINI_API_KEY is set. Error: {e}")
    exit(1)

MAX_WORKERS = len(API_TOKENS)
token_queue = queue.Queue()
for t in API_TOKENS:
    token_queue.put(t)

# ==============================================================================
# GEMINI EMBEDDING (Cache-Free & Rate-Limit Safe)
# ==============================================================================
def get_gemini_embeddings(texts):
    """Fetches embeddings directly from Gemini with exponential backoff."""
    embeddings = [None] * len(texts)

    def _fetch_with_backoff(j, text):
        max_retries = 6
        for attempt in range(max_retries):
            try:
                res = gemini_client.models.embed_content(
                    model=EMBEDDING_MODEL,
                    contents=text,
                    config=types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY")
                )
                return j, list(res.embeddings[0].values)
            except Exception as e:
                err_str = str(e).lower()
                if "429" in err_str or "exhausted" in err_str or "quota" in err_str:
                    time.sleep((2 ** attempt) * 2)
                else:
                    raise e
        raise Exception("Max retries exceeded for Gemini API.")

    # Max 5 workers to respect Gemini limits
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        futures = [executor.submit(_fetch_with_backoff, j, text) for j, text in enumerate(texts)]
        for future in concurrent.futures.as_completed(futures):
            j, emb_values = future.result()
            embeddings[j] = emb_values

    return np.array(embeddings, dtype=np.float32)


# ==============================================================================
# HUMAN DATA LOADING
# ==============================================================================
HUMAN_TEXT_DIR = "./optuna_human"
HUMAN_TEXTS = []

print(f"Loading human reference texts from {HUMAN_TEXT_DIR}...")
print("For human reference texts, cutoff is 5000 characters.")
print(f"Each file should be the exact same token count as 1024 + CONTEXT_PREFIX token count.")
if not os.path.exists(HUMAN_TEXT_DIR):
    os.makedirs(HUMAN_TEXT_DIR)
    print(f"[!] Created directory {HUMAN_TEXT_DIR}. Add .txt files and restart.")
    exit(1)

txt_files = sorted(glob.glob(os.path.join(HUMAN_TEXT_DIR, "*.txt")))
for filepath in txt_files:
    with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
        text = f.read().strip()
        if not text:
            continue
        HUMAN_TEXTS.append(text[:5000])

print(f" -> Successfully loaded {len(HUMAN_TEXTS)} human files.")

human_dataset_string = "".join(sorted(HUMAN_TEXTS)).encode("utf-8")
human_hash = hashlib.sha256(human_dataset_string).hexdigest()[:8]

print("Pre-computing Gemini embeddings for human baseline...")
HUMAN_EMBEDDINGS = get_gemini_embeddings(HUMAN_TEXTS)
HUMAN_EMBEDDINGS = HUMAN_EMBEDDINGS / np.linalg.norm(
    HUMAN_EMBEDDINGS, axis=1, keepdims=True
)
print(f" -> Extracted {HUMAN_EMBEDDINGS.shape[0]} embeddings.\n")


# ==============================================================================
# GENERATION WORKER
# ==============================================================================
def generate_sample(params, prompt_tokens, cancel_event: threading.Event):
    current_token = token_queue.get()
    try:
        payload = {
            "model": NAI_MODEL,
            "logit_bias": LOGIT_BIAS,
            "max_tokens": 1024,
            "temperature": params["temp"],
            "top_p": params["top_p"],
            "min_p": params["min_p"],
            "frequency_penalty": params["freq"],
            "presence_penalty": params["pres"],
            "top_k": params["top_k"],
            "prompt": prompt_tokens,
            "stream": True,
        }
        headers = {"Authorization": f"Bearer {current_token}"}

        generated_chars = []
        generated_tokens = []
        generated = 0

        while not cancel_event.is_set():
            try:
                payload["prompt"] = prompt_tokens + generated_tokens
                with requests.post(
                    API_URL, json=payload, headers=headers,
                    timeout=120, stream=True,
                ) as resp:

                    if resp.status_code != 200:
                        if resp.status_code in (400, 401, 402, 403, 404):
                            return None
                        if cancel_event.wait(timeout=2):
                            return None
                        continue  # transient — retry

                    for raw in resp.iter_lines(decode_unicode=True):
                        if cancel_event.is_set():
                            return None  # `with` closes socket on exit

                        if not raw:
                            continue
                        if raw.startswith("data:"):
                            raw = raw[5:].lstrip()
                        if raw == "[DONE]":
                            break

                        try:
                            char = json.loads(raw)["choices"][0].get("text", "")
                            token = json.loads(raw)["choices"][0].get("token_ids", [])
                        except (json.JSONDecodeError, KeyError, IndexError):
                            continue

                        if not char and not token:
                            continue

                        generated_chars.append(char)
                        generated_tokens.extend(token)
                        generated += 1
                        print(".", end="", flush=True)
                        if generated >= 1024:
                            print()
                            return CONTEXT_PREFIX + "".join(generated_chars)

                    # Stream ended before 1024 — too short, retry
            except Exception:
                if cancel_event.wait(timeout=2):
                    return None
                continue

        return None
    finally:
        token_queue.put(current_token)


# ==============================================================================
# OPTIMIZATION LOOP
# ==============================================================================
def objective(trial):
    params = {
        "temp":  trial.suggest_float("temp", 0.1, 2),
        "top_p": trial.suggest_float("top_p", 0.05, 1.0),
        "min_p": trial.suggest_float("min_p", 0.0, 1.0),
        "freq":  trial.suggest_float("freq", 0.0, 1.0),
        "pres":  trial.suggest_float("pres", 0.0, 2.0),
        "top_k": trial.suggest_int("top_k", 1, 151364),
    }

    print(f"\n---------------------- TRIAL #{trial.number} ----------------------")
    print(params)

    completed_count = 0
    geometric_mean    = 0.0
    failed_count    = 0
    running_log_mean = 0.0 

    cancel_event = threading.Event()

    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures   = {}
        submitted = 0
        pruned    = False

        def submit_next():
            """Submit the next eval_id if any remain. Returns True if submitted."""
            nonlocal submitted
            if submitted >= MAX_SAMPLES_PER_TRIAL:
                return False
            fut = executor.submit(generate_sample, params, PROMPT_TOKENS, cancel_event)
            futures[fut] = submitted
            submitted += 1
            return True

        # Prime the pool
        for _ in range(min(MAX_WORKERS, MAX_SAMPLES_PER_TRIAL)):
            submit_next()

        while futures and not pruned:
            done, _ = concurrent.futures.wait(
                futures, return_when=concurrent.futures.FIRST_COMPLETED
            )

            for f in done:
                eval_id = futures.pop(f)

                # 1. Get the worker's result, surviving worker-side exceptions.
                try:
                    txt = f.result()
                except Exception as e:
                    print(f"   [!] Worker raised for eval {eval_id}: {e!r}")
                    txt = None

                # 2. Score it, surviving embedding/reporting exceptions.
                if txt:
                    try:
                        emb = get_gemini_embeddings([txt])[0]
                        emb_n = emb / np.linalg.norm(emb)
                        sims = HUMAN_EMBEDDINGS @ emb_n           # shape (3,), each in [-1, 1]

                        cos_sim = float(sims.max())
                        completed_count += 1

                        safe_sim = max(cos_sim, 1e-8) 
                        running_log_mean += (math.log(safe_sim) - running_log_mean) / completed_count
                        geometric_mean = math.exp(running_log_mean)
                        step = completed_count - 1

                        print(f"TRIAL #{trial.number} | SAMPLE #{step}/{MAX_SAMPLES_PER_TRIAL} | MAX COSINE: {cos_sim:.16f} | GEOMETRIC MEAN: {geometric_mean:.16f}")

                        trial.report(geometric_mean, step)
    
                        if trial.should_prune():
                            print(f" [!] Pruned at step {step}. Signalling workers to stop.")
                            cancel_event.set()
                            pruned = True
                    except Exception as e:
                        failed_count += 1
                        print(f"   [!] Scoring failed for eval {eval_id}: {e!r}")
                else:
                    failed_count += 1
                    print(f"   [!] Empty/missing result for eval {eval_id}; topping up queue.")

                # 3. ALWAYS try to keep the pool full, regardless of success/failure.
                if not pruned:
                    submit_next()

        # Drain cancelled futures cleanly if we pruned mid-flight
        if pruned:
            for f in futures:
                f.cancel()

    if pruned:
        raise optuna.TrialPruned()

    # Guard: if every sample failed, don't feed 0.0 to Optuna as a real score.
    if completed_count == 0:
        raise RuntimeError(f"Trial {trial.number}: all {failed_count} samples failed")

    return geometric_mean


# ==============================================================================
# ENTRY POINT
# ==============================================================================
if __name__ == "__main__":
    STUDY_NAME = f"hail_mary_preset_s{MAX_SAMPLES_PER_TRIAL}_h{human_hash}_cosine"

    print(f"\nStarting Optuna Optimization Pipeline")
    print(f"Scoring Mode: COSINE")
    print(f"Current Target: {STUDY_NAME}\n")

    study = optuna.create_study(
        study_name=STUDY_NAME,
        storage="sqlite:///optuna_xialong_1.db",
        load_if_exists=True,
        direction="maximize",
        sampler=optuna.samplers.TPESampler(),
        pruner=optuna.pruners.HyperbandPruner()
    )

    valid_states = [optuna.trial.TrialState.COMPLETE, optuna.trial.TrialState.RUNNING]
    existing_trials = [t.params for t in study.get_trials() if t.state in valid_states]
    print(f"Found {len(existing_trials)} finished or running trials.")
    if len(existing_trials) == 0:
        study.enqueue_trial(INITIAL_PRESET)
        print(" -> Enqueued initial trial.")

    print("==================================================\n")
    study.optimize(objective)