bookapp/story/bible_tracker.py

import json
from core import utils
from ai import models as ai_models


def merge_selected_changes(original, draft, selected_keys):
    def sort_key(k):
        return [int(p) if p.isdigit() else p for p in k.split('.')]
    selected_keys.sort(key=sort_key)

    for key in selected_keys:
        parts = key.split('.')

        if parts[0] == 'meta' and len(parts) == 2:
            field = parts[1]
            if field == 'tone':
                original['project_metadata']['style']['tone'] = draft['project_metadata']['style']['tone']
            elif field in original['project_metadata']:
                original['project_metadata'][field] = draft['project_metadata'][field]

        elif parts[0] == 'char' and len(parts) >= 2:
            idx = int(parts[1])
            if idx < len(draft['characters']):
                if idx < len(original['characters']):
                    original['characters'][idx] = draft['characters'][idx]
                else:
                    original['characters'].append(draft['characters'][idx])

        elif parts[0] == 'book' and len(parts) >= 2:
            book_num = int(parts[1])
            orig_book = next((b for b in original['books'] if b['book_number'] == book_num), None)
            draft_book = next((b for b in draft['books'] if b['book_number'] == book_num), None)

            if draft_book:
                if not orig_book:
                    original['books'].append(draft_book)
                    original['books'].sort(key=lambda x: x.get('book_number', 999))
                    continue

                if len(parts) == 2:
                    orig_book['title'] = draft_book['title']
                    orig_book['manual_instruction'] = draft_book['manual_instruction']

                elif len(parts) == 4 and parts[2] == 'beat':
                    beat_idx = int(parts[3])
                    if beat_idx < len(draft_book['plot_beats']):
                        while len(orig_book['plot_beats']) <= beat_idx:
                            orig_book['plot_beats'].append("")
                        orig_book['plot_beats'][beat_idx] = draft_book['plot_beats'][beat_idx]
    return original


def filter_characters(chars):
    blacklist = ['name', 'character name', 'role', 'protagonist', 'antagonist', 'love interest', 'unknown', 'tbd', 'todo', 'hero', 'villain', 'main character', 'side character']
    return [c for c in chars if c.get('name') and c.get('name').lower().strip() not in blacklist]


def update_tracking(folder, chapter_num, chapter_text, current_tracking):
    utils.log("TRACKER", f"Updating world state & character visuals for Ch {chapter_num}...")

    prompt = f"""
    ROLE: Continuity Tracker
    TASK: Update the Story Bible based on the new chapter.

    INPUT_TRACKING:
    {json.dumps(current_tracking)}

    NEW_TEXT:
    {chapter_text[:20000]}

    OPERATIONS:
    1. EVENTS: Append 1-3 key plot points to 'events'.
    2. CHARACTERS: Update 'descriptors', 'likes_dislikes', 'speech_style', 'last_worn', 'major_events', 'current_location', 'time_of_day', 'held_items'.
       - "descriptors": List of strings. Add PERMANENT physical traits (height, hair, eyes), specific items (jewelry, weapons). Avoid duplicates.
       - "likes_dislikes": List of strings. Add specific preferences, likes, or dislikes mentioned (e.g., "Hates coffee", "Loves jazz").
       - "speech_style": String. Describe how they speak (e.g. "Formal, no contractions", "Uses slang", "Stutters", "Short sentences").
       - "last_worn": String. Update if specific clothing is described. IMPORTANT: If a significant time jump occurred (e.g. next day) and no new clothing is described, reset this to "Unknown".
       - "major_events": List of strings. Log significant life-altering events occurring in THIS chapter (e.g. "Lost an arm", "Married", "Betrayed by X").
       - "current_location": String. The character's physical location at the END of this chapter (e.g., "The King's Throne Room", "Aboard the Nighthawk ship"). Update whenever the character moves.
       - "time_of_day": String. The approximate time of day at the END of this chapter (e.g., "Dawn", "Late afternoon", "Midnight"). Reset to "Unknown" if unclear.
       - "held_items": List of strings. Items the character is actively carrying or holding at chapter end (e.g., "Iron sword", "Stolen ledger"). Remove items they have dropped or given away.
    3. WARNINGS: Append new 'content_warnings'.

    OUTPUT_FORMAT (JSON): Return the updated tracking object structure.
    """
    try:
        response = ai_models.model_logic.generate_content(prompt)
        utils.log_usage(folder, ai_models.model_logic.name, response.usage_metadata)
        new_data = json.loads(utils.clean_json(response.text))
        return new_data
    except Exception as e:
        utils.log("TRACKER", f"Failed to update tracking: {e}")
        return current_tracking


def update_lore_index(folder, chapter_text, current_lore):
    """Extract canonical descriptions of locations and key items from a chapter
    and merge them into the lore index dict. Returns the updated lore dict."""
    utils.log("TRACKER", "Updating lore index from chapter...")
    prompt = f"""
    ROLE: Lore Keeper
    TASK: Extract canonical descriptions of locations and key items from this chapter.

    EXISTING_LORE:
    {json.dumps(current_lore)}

    CHAPTER_TEXT:
    {chapter_text[:15000]}

    INSTRUCTIONS:
    1. For each LOCATION mentioned: provide a 1-2 sentence canonical description (appearance, atmosphere, notable features).
    2. For each KEY ITEM or ARTIFACT mentioned: provide a 1-2 sentence canonical description (appearance, properties, significance).
    3. Do NOT add characters — only physical places and objects.
    4. If an entry already exists in EXISTING_LORE, update or preserve it — do not duplicate.
    5. Use the exact name as the key (e.g., "The Thornwood Inn", "The Sunstone Amulet").
    6. Only include entries that have meaningful descriptive detail in the chapter text.

    OUTPUT_FORMAT (JSON): {{"LocationOrItemName": "Description.", ...}}
    """
    try:
        response = ai_models.model_logic.generate_content(prompt)
        utils.log_usage(folder, ai_models.model_logic.name, response.usage_metadata)
        new_entries = json.loads(utils.clean_json(response.text))
        if isinstance(new_entries, dict):
            current_lore.update(new_entries)
        return current_lore
    except Exception as e:
        utils.log("TRACKER", f"Lore index update failed: {e}")
        return current_lore


def harvest_metadata(bp, folder, full_manuscript):
    utils.log("HARVESTER", "Scanning for new characters...")
    full_text = "\n".join([c.get('content', '') for c in full_manuscript])[:500000]

    prompt = f"""
    ROLE: Data Extractor
    TASK: Identify NEW significant characters.

    INPUT_TEXT:
    {full_text}

    KNOWN_CHARACTERS: {json.dumps(bp['characters'])}

    OUTPUT_FORMAT (JSON): {{ "new_characters": [{{ "name": "String", "role": "String", "description": "String" }}] }}
    """
    try:
        response = ai_models.model_logic.generate_content(prompt)
        utils.log_usage(folder, ai_models.model_logic.name, response.usage_metadata)
        new_chars = json.loads(utils.clean_json(response.text)).get('new_characters', [])
        if new_chars:
            valid_chars = filter_characters(new_chars)
            if valid_chars:
                utils.log("HARVESTER", f"Found {len(valid_chars)} new chars.")
                bp['characters'].extend(valid_chars)
    except: pass
    return bp


def refine_bible(bible, instruction, folder):
    utils.log("SYSTEM", f"Refining Bible with instruction: {instruction}")
    prompt = f"""
    ROLE: Senior Developmental Editor
    TASK: Update the Bible JSON based on instruction.

    INPUT_DATA:
    - CURRENT_JSON: {json.dumps(bible)}
    - INSTRUCTION: {instruction}

    CONSTRAINTS:
    - Maintain valid JSON structure.
    - Ensure consistency.

    OUTPUT_FORMAT (JSON): The full updated Bible JSON object.
    """
    try:
        response = ai_models.model_logic.generate_content(prompt)
        utils.log_usage(folder, ai_models.model_logic.name, response.usage_metadata)
        new_data = json.loads(utils.clean_json(response.text))
        return new_data
    except Exception as e:
        utils.log("SYSTEM", f"Refinement failed: {e}")
        return None