bookapp/modules/utils.py

import os
import json
import datetime
import time
import config
import threading
import re

SAFETY_SETTINGS = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
]

# Thread-local storage for logging context
_log_context = threading.local()

# Cache for dynamic pricing from AI model selection
PRICING_CACHE = {}

def set_log_file(filepath):
    _log_context.log_file = filepath

def set_log_callback(callback):
    _log_context.callback = callback

def set_progress_callback(callback):
    _log_context.progress_callback = callback

def update_progress(percent):
    if getattr(_log_context, 'progress_callback', None):
        try: _log_context.progress_callback(percent)
        except: pass

def clean_json(text):
    text = text.replace("```json", "").replace("```", "").strip()
    # Robust extraction: find first { or [ and last } or ]
    start_obj = text.find('{')
    start_arr = text.find('[')
    if start_obj == -1 and start_arr == -1: return text
    if start_obj != -1 and (start_arr == -1 or start_obj < start_arr):
        return text[start_obj:text.rfind('}')+1]
    else:
        return text[start_arr:text.rfind(']')+1]

def sanitize_filename(name):
    """Sanitizes a string to be safe for filenames."""
    if not name: return "Untitled"
    safe = "".join([c for c in name if c.isalnum() or c=='_']).replace(" ", "_")
    return safe if safe else "Untitled"

def chapter_sort_key(ch):
    """Sort key for chapters handling integers, strings, Prologue, and Epilogue."""
    num = ch.get('num', 0)
    if isinstance(num, int): return num
    if isinstance(num, str) and num.isdigit(): return int(num)
    s = str(num).lower().strip()
    if 'prologue' in s: return -1
    if 'epilogue' in s: return 9999
    return 999

def get_sorted_book_folders(run_dir):
    """Returns a list of book folder names in a run directory, sorted numerically."""
    if not os.path.exists(run_dir): return []
    subdirs = [d for d in os.listdir(run_dir) if os.path.isdir(os.path.join(run_dir, d)) and d.startswith("Book_")]
    def sort_key(d):
        parts = d.split('_')
        if len(parts) > 1 and parts[1].isdigit(): return int(parts[1])
        return 0
    return sorted(subdirs, key=sort_key)

# --- SHARED UTILS ---
def log_banner(phase, title):
    """Log a visually distinct phase separator line."""
    log(phase, f"{'─' * 18} {title} {'─' * 18}")

def log(phase, msg):
    timestamp = datetime.datetime.now().strftime('%H:%M:%S')
    line = f"[{timestamp}] {phase:<15} | {msg}"
    print(line)
    
    # Write to thread-specific log file if set
    if getattr(_log_context, 'log_file', None):
        with open(_log_context.log_file, "a", encoding="utf-8") as f:
            f.write(line + "\n")
            
    # Trigger callback if set (e.g. for Database logging)
    if getattr(_log_context, 'callback', None):
        try: _log_context.callback(phase, msg)
        except: pass

def load_json(path): 
    return json.load(open(path, 'r')) if os.path.exists(path) else None

def create_default_personas():
    # Initialize empty personas file if it doesn't exist
    if not os.path.exists(config.PERSONAS_DIR): os.makedirs(config.PERSONAS_DIR)
    if not os.path.exists(config.PERSONAS_FILE):
        try:
            with open(config.PERSONAS_FILE, 'w') as f: json.dump({}, f, indent=2)
        except: pass

def get_length_presets():
    """Returns a dict mapping Label -> Settings for use in main.py"""
    presets = {}
    for k, v in config.LENGTH_DEFINITIONS.items():
        presets[v['label']] = v
    return presets

def log_image_attempt(folder, img_type, prompt, filename, status, error=None, score=None, critique=None):
    log_path = os.path.join(folder, "image_log.json")
    entry = {
        "timestamp": int(time.time()),
        "type": img_type,
        "prompt": prompt,
        "filename": filename,
        "status": status,
        "error": str(error) if error else None,
        "score": score,
        "critique": critique
    }
    data = []
    if os.path.exists(log_path):
        try:
            with open(log_path, 'r') as f: data = json.load(f)
        except:
            pass
    data.append(entry)
    with open(log_path, 'w') as f: json.dump(data, f, indent=2)

def get_run_folder(base_name):
    if not os.path.exists(base_name): os.makedirs(base_name)
    runs = [d for d in os.listdir(base_name) if d.startswith("run_")]
    next_num = max([int(r.split("_")[1]) for r in runs if r.split("_")[1].isdigit()] + [0]) + 1
    folder = os.path.join(base_name, f"run_{next_num}")
    os.makedirs(folder)
    return folder

def get_latest_run_folder(base_name):
    if not os.path.exists(base_name): return None
    runs = [d for d in os.listdir(base_name) if d.startswith("run_")]
    if not runs: return None
    runs.sort(key=lambda x: int(x.split('_')[1]) if x.split('_')[1].isdigit() else 0)
    return os.path.join(base_name, runs[-1])

def update_pricing(model_name, cost_str):
    """Parses cost string from AI selection and updates cache."""
    if not model_name or not cost_str or cost_str == 'N/A': return
    
    try:
        # Look for patterns like "$0.075 Input" or "$3.50/1M"
        # Default to 0.0
        in_cost = 0.0
        out_cost = 0.0
        
        # Extract all float-like numbers following a $ sign
        prices = re.findall(r'(?:\$|USD)\s*([0-9]+\.?[0-9]*)', cost_str, re.IGNORECASE)
        
        if len(prices) >= 2:
            in_cost = float(prices[0])
            out_cost = float(prices[1])
        elif len(prices) == 1:
            in_cost = float(prices[0])
            out_cost = in_cost * 3 # Rough heuristic if only one price provided
            
        if in_cost > 0:
            PRICING_CACHE[model_name] = {"input": in_cost, "output": out_cost}
            # log("SYSTEM", f"Updated pricing for {model_name}: In=${in_cost} | Out=${out_cost}")
    except:
        pass

def calculate_cost(model_label, input_tokens, output_tokens, image_count=0):
    cost = 0.0
    m = model_label.lower()
    
    # Check dynamic cache first
    if model_label in PRICING_CACHE:
        rates = PRICING_CACHE[model_label]
        cost = (input_tokens / 1_000_000 * rates['input']) + (output_tokens / 1_000_000 * rates['output'])
    elif 'imagen' in m or image_count > 0:
        cost = (image_count * 0.04)
    else:
        # Fallbacks
        if 'flash' in m:
            cost = (input_tokens / 1_000_000 * 0.075) + (output_tokens / 1_000_000 * 0.30)
        elif 'pro' in m or 'logic' in m:
            cost = (input_tokens / 1_000_000 * 3.50) + (output_tokens / 1_000_000 * 10.50)
            
    return round(cost, 6)

def log_usage(folder, model_label, usage_metadata=None, image_count=0):
    if not folder or not os.path.exists(folder): return
    
    log_path = os.path.join(folder, "usage_log.json")
    
    input_tokens = 0
    output_tokens = 0
    
    if usage_metadata:
        try:
            input_tokens = usage_metadata.prompt_token_count
            output_tokens = usage_metadata.candidates_token_count
        except: pass

    # Calculate Cost
    cost = calculate_cost(model_label, input_tokens, output_tokens, image_count)

    entry = {
        "timestamp": int(time.time()),
        "date": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        "model": model_label,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "images": image_count,
        "cost": round(cost, 6)
    }

    data = {"log": [], "totals": {"input_tokens": 0, "output_tokens": 0, "images": 0, "est_cost_usd": 0.0}}
    
    if os.path.exists(log_path):
        try:
            loaded = json.load(open(log_path, 'r'))
            if isinstance(loaded, list): data["log"] = loaded
            elif isinstance(loaded, dict): data = loaded
        except: pass
        
    data["log"].append(entry)
    
    # Recalculate totals
    t_in = sum(x.get('input_tokens', 0) for x in data["log"])
    t_out = sum(x.get('output_tokens', 0) for x in data["log"])
    t_img = sum(x.get('images', 0) for x in data["log"])
    
    total_cost = 0.0
    for x in data["log"]:
        if 'cost' in x:
            total_cost += x['cost']
        else:
            # Fallback calculation for old logs without explicit cost field
            c = 0.0
            mx = x.get('model', '').lower()
            ix = x.get('input_tokens', 0)
            ox = x.get('output_tokens', 0)
            imgx = x.get('images', 0)
            
            if 'flash' in mx: c = (ix / 1_000_000 * 0.075) + (ox / 1_000_000 * 0.30)
            elif 'pro' in mx or 'logic' in mx: c = (ix / 1_000_000 * 3.50) + (ox / 1_000_000 * 10.50)
            elif 'imagen' in mx or imgx > 0: c = (imgx * 0.04)
            total_cost += c
            
    data["totals"] = {
        "input_tokens": t_in,
        "output_tokens": t_out,
        "images": t_img,
        "est_cost_usd": round(total_cost, 4)
    }
    
    with open(log_path, 'w') as f: json.dump(data, f, indent=2)