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:
            try:
                idx = int(parts[1])
            except (ValueError, IndexError):
                utils.log("SYSTEM", f"⚠️ Skipping malformed bible merge key: '{key}'")
                continue
            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:
            try:
                book_num = int(parts[1])
            except (ValueError, IndexError):
                utils.log("SYSTEM", f"⚠️ Skipping malformed bible merge key: '{key}'")
                continue
            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':
                    try:
                        beat_idx = int(parts[3])
                    except (ValueError, IndexError):
                        utils.log("SYSTEM", f"⚠️ Skipping malformed beat merge key: '{key}'")
                        continue
                    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 merge_tracking_to_bible(bible, tracking):
    """Merge dynamic tracking state back into the bible dict.

    Makes bible.json the single persistent source of truth by updating
    character data and lore from the in-memory tracking object.
    Returns the modified bible dict.
    """
    for name, data in tracking.get('characters', {}).items():
        matched = False
        for char in bible.get('characters', []):
            if char.get('name') == name:
                char.update(data)
                matched = True
                break
        if not matched:
            utils.log("TRACKER", f"  -> Character '{name}' in tracking not found in bible. Skipping.")

    if 'lore' not in bible:
        bible['lore'] = {}
    bible['lore'].update(tracking.get('lore', {}))

    return bible


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 Exception as e:
        utils.log("HARVESTER", f"⚠️ Metadata harvest failed: {e}")
    return bp


def get_chapter_neighbours(manuscript, current_num):
    """Return (prev_num, next_num) chapter numbers adjacent to current_num.

    manuscript: list of chapter dicts each with a 'num' key.
    Returns None for prev/next when at the boundary.
    """
    nums = sorted({ch.get('num') for ch in manuscript if ch.get('num') is not None})
    if current_num not in nums:
        return None, None
    idx = nums.index(current_num)
    prev_num = nums[idx - 1] if idx > 0 else None
    next_num = nums[idx + 1] if idx < len(nums) - 1 else None
    return prev_num, next_num


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