character-browser/routes/regenerate.py

import json
import os
import logging

from flask import current_app
from sqlalchemy.orm.attributes import flag_modified

from models import db, Character, Outfit, Action, Style, Scene, Detailer, Look, Checkpoint
from services.llm import load_prompt, call_llm
from services.sync import _sync_nsfw_from_tags
from services.job_queue import _enqueue_task

logger = logging.getLogger('gaze')

# Map category string to (model class, id_field, config_dir_key)
_CATEGORY_MAP = {
    'characters': (Character, 'character_id', 'CHARACTERS_DIR'),
    'outfits':    (Outfit,    'outfit_id',    'CLOTHING_DIR'),
    'actions':    (Action,    'action_id',    'ACTIONS_DIR'),
    'styles':     (Style,     'style_id',     'STYLES_DIR'),
    'scenes':     (Scene,     'scene_id',     'SCENES_DIR'),
    'detailers':  (Detailer,  'detailer_id',  'DETAILERS_DIR'),
    'looks':      (Look,      'look_id',      'LOOKS_DIR'),
}

# Fields to preserve from the original data (never overwritten by LLM output)
_PRESERVE_KEYS = {
    'lora', 'participants', 'suppress_wardrobe',
    'character_id', 'character_name',
    'outfit_id', 'outfit_name',
    'action_id', 'action_name',
    'style_id', 'style_name',
    'scene_id', 'scene_name',
    'detailer_id', 'detailer_name',
    'look_id', 'look_name',
}


def register_routes(app):

    @app.route('/api/<category>/<path:slug>/regenerate_tags', methods=['POST'])
    def regenerate_tags(category, slug):
        if category not in _CATEGORY_MAP:
            return {'error': f'Unknown category: {category}'}, 400

        model_class, id_field, dir_key = _CATEGORY_MAP[category]
        entity = model_class.query.filter_by(slug=slug).first()
        if not entity:
            return {'error': 'Not found'}, 404

        system_prompt = load_prompt('regenerate_tags_system.txt')
        if not system_prompt:
            return {'error': 'Regenerate tags system prompt not found'}, 500

        original_data = entity.data.copy()

        try:
            prompt = (
                f"Regenerate the prompt tags for this {category.rstrip('s')} resource.\n"
                f"Current JSON:\n{json.dumps(original_data, indent=2)}"
            )
            llm_response = call_llm(prompt, system_prompt)
            clean_json = llm_response.replace('```json', '').replace('```', '').strip()
            new_data = json.loads(clean_json)
        except Exception as e:
            logger.exception(f"Regenerate tags LLM error for {category}/{slug}")
            return {'error': f'LLM error: {str(e)}'}, 500

        # Preserve protected fields from original
        for key in _PRESERVE_KEYS:
            if key in original_data:
                new_data[key] = original_data[key]

        # Update DB
        entity.data = new_data
        flag_modified(entity, 'data')
        _sync_nsfw_from_tags(entity, new_data)
        db.session.commit()

        # Write back to JSON file
        if entity.filename:
            file_path = os.path.join(current_app.config[dir_key], entity.filename)
            with open(file_path, 'w') as f:
                json.dump(new_data, f, indent=2)

        return {'success': True, 'data': new_data}

    @app.route('/admin/migrate_tags', methods=['POST'])
    def migrate_tags():
        """One-time migration: convert old list-format tags to new dict format."""
        migrated = 0
        for category, (model_class, id_field, dir_key) in _CATEGORY_MAP.items():
            entities = model_class.query.all()
            for entity in entities:
                tags = entity.data.get('tags')
                if isinstance(tags, list) or tags is None:
                    new_data = entity.data.copy()
                    new_data['tags'] = {'nsfw': False}
                    entity.data = new_data
                    flag_modified(entity, 'data')

                    # Write back to JSON file
                    if entity.filename:
                        file_path = os.path.join(current_app.config[dir_key], entity.filename)
                        try:
                            with open(file_path, 'w') as f:
                                json.dump(new_data, f, indent=2)
                        except Exception as e:
                            logger.warning(f"Could not write {file_path}: {e}")

                    migrated += 1

        # Also handle checkpoints
        for ckpt in Checkpoint.query.all():
            data = ckpt.data or {}
            tags = data.get('tags')
            if isinstance(tags, list) or tags is None:
                new_data = data.copy()
                new_data['tags'] = {'nsfw': False}
                ckpt.data = new_data
                flag_modified(ckpt, 'data')
                migrated += 1

        db.session.commit()
        logger.info(f"Migrated {migrated} resources from list tags to dict tags")
        return {'success': True, 'migrated': migrated}

    def _make_regen_task(category, slug, name, system_prompt):
        """Factory: create a tag regeneration task function for one entity."""
        def task_fn(job):
            model_class, id_field, dir_key = _CATEGORY_MAP[category]
            entity = model_class.query.filter_by(slug=slug).first()
            if not entity:
                raise Exception(f'{category}/{slug} not found')

            original_data = entity.data.copy()
            prompt = (
                f"Regenerate the prompt tags for this {category.rstrip('s')} resource.\n"
                f"Current JSON:\n{json.dumps(original_data, indent=2)}"
            )
            llm_response = call_llm(prompt, system_prompt)
            clean_json = llm_response.replace('```json', '').replace('```', '').strip()
            new_data = json.loads(clean_json)

            for key in _PRESERVE_KEYS:
                if key in original_data:
                    new_data[key] = original_data[key]

            entity.data = new_data
            flag_modified(entity, 'data')
            _sync_nsfw_from_tags(entity, new_data)
            db.session.commit()

            if entity.filename:
                file_path = os.path.join(current_app.config[dir_key], entity.filename)
                with open(file_path, 'w') as f:
                    json.dump(new_data, f, indent=2)

            job['result'] = {'entity': name, 'status': 'updated'}
        return task_fn

    @app.route('/admin/bulk_regenerate_tags/<category>', methods=['POST'])
    def bulk_regenerate_tags_category(category):
        """Queue LLM tag regeneration for all resources in a single category."""
        if category not in _CATEGORY_MAP:
            return {'error': f'Unknown category: {category}'}, 400

        system_prompt = load_prompt('regenerate_tags_system.txt')
        if not system_prompt:
            return {'error': 'Regenerate tags system prompt not found'}, 500

        model_class, id_field, dir_key = _CATEGORY_MAP[category]
        entities = model_class.query.all()
        job_ids = []

        for entity in entities:
            job = _enqueue_task(
                f"Regen tags: {entity.name} ({category})",
                _make_regen_task(category, entity.slug, entity.name, system_prompt)
            )
            job_ids.append(job['id'])

        return {'success': True, 'queued': len(job_ids), 'job_ids': job_ids}

    @app.route('/admin/bulk_regenerate_tags', methods=['POST'])
    def bulk_regenerate_tags():
        """Queue LLM tag regeneration for all resources across all categories."""
        system_prompt = load_prompt('regenerate_tags_system.txt')
        if not system_prompt:
            return {'error': 'Regenerate tags system prompt not found'}, 500

        job_ids = []
        for category, (model_class, id_field, dir_key) in _CATEGORY_MAP.items():
            entities = model_class.query.all()
            for entity in entities:
                job = _enqueue_task(
                    f"Regen tags: {entity.name} ({category})",
                    _make_regen_task(category, entity.slug, entity.name, system_prompt)
                )
                job_ids.append(job['id'])

        return {'success': True, 'queued': len(job_ids), 'job_ids': job_ids}