import json
import os
import logging

from flask import render_template, request, current_app
from models import (
    db, Character, Outfit, Action, Style, Scene, Detailer, Look, Checkpoint, Preset,
)

logger = logging.getLogger('gaze')


GALLERY_CATEGORIES = ['characters', 'actions', 'outfits', 'scenes', 'styles', 'detailers', 'checkpoints', 'looks', 'presets', 'generator']

_MODEL_MAP = {
    'characters': Character,
    'actions': Action,
    'outfits': Outfit,
    'scenes': Scene,
    'styles': Style,
    'detailers': Detailer,
    'checkpoints': Checkpoint,
    'looks': Look,
    'presets': Preset,
    'generator': Preset,
}

# Maps xref_category param names to sidecar JSON keys
_XREF_KEY_MAP = {
    'character': 'character_slug',
    'outfit': 'outfit_slug',
    'action': 'action_slug',
    'style': 'style_slug',
    'scene': 'scene_slug',
    'detailer': 'detailer_slug',
    'look': 'look_slug',
    'preset': 'preset_slug',
}


def register_routes(app):

    def _read_sidecar(upload_folder, image_path):
        """Read JSON sidecar for an image. Returns dict or None."""
        sidecar = image_path.rsplit('.', 1)[0] + '.json'
        sidecar_path = os.path.join(upload_folder, sidecar)
        try:
            with open(sidecar_path) as f:
                return json.load(f)
        except (OSError, json.JSONDecodeError):
            return None

    def _scan_gallery_images(category_filter='all', slug_filter=''):
        """Return sorted list of image dicts from the uploads directory."""
        upload_folder = app.config['UPLOAD_FOLDER']
        images = []
        cats = GALLERY_CATEGORIES if category_filter == 'all' else [category_filter]

        for cat in cats:
            cat_folder = os.path.join(upload_folder, cat)
            if not os.path.isdir(cat_folder):
                continue
            try:
                slugs = os.listdir(cat_folder)
            except OSError:
                continue
            for item_slug in slugs:
                if slug_filter and slug_filter != item_slug:
                    continue
                item_folder = os.path.join(cat_folder, item_slug)
                if not os.path.isdir(item_folder):
                    continue
                try:
                    files = os.listdir(item_folder)
                except OSError:
                    continue
                for filename in files:
                    if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
                        continue
                    try:
                        ts = int(filename.replace('gen_', '').rsplit('.', 1)[0])
                    except ValueError:
                        ts = 0
                    images.append({
                        'path': f"{cat}/{item_slug}/{filename}",
                        'category': cat,
                        'slug': item_slug,
                        'filename': filename,
                        'timestamp': ts,
                    })

        images.sort(key=lambda x: x['timestamp'], reverse=True)
        return images

    def _enrich_with_names(images):
        """Add item_name field to each image dict, querying DB once per category."""
        by_cat = {}
        for img in images:
            by_cat.setdefault(img['category'], set()).add(img['slug'])

        name_map = {}
        for cat, slugs in by_cat.items():
            Model = _MODEL_MAP.get(cat)
            if not Model:
                continue
            items = Model.query.filter(Model.slug.in_(slugs)).with_entities(Model.slug, Model.name).all()
            for slug, name in items:
                name_map[(cat, slug)] = name

        for img in images:
            img['item_name'] = name_map.get((img['category'], img['slug']), img['slug'])
        return images

    def _parse_comfy_png_metadata(image_path):
        """Read ComfyUI generation metadata from a PNG's tEXt 'prompt' chunk.

        Returns a dict with keys: positive, negative, checkpoint, loras,
        seed, steps, cfg, sampler, scheduler.  Any missing field is None/[].
        """
        from PIL import Image as PilImage

        result = {
            'positive': None,
            'negative': None,
            'checkpoint': None,
            'loras': [],       # list of {name, strength}
            'seed': None,
            'steps': None,
            'cfg': None,
            'sampler': None,
            'scheduler': None,
        }

        try:
            with PilImage.open(image_path) as im:
                raw = im.info.get('prompt')
            if not raw:
                return result
            nodes = json.loads(raw)
        except Exception:
            return result

        for node in nodes.values():
            ct = node.get('class_type', '')
            inp = node.get('inputs', {})

            if ct == 'KSampler':
                result['seed']      = inp.get('seed')
                result['steps']     = inp.get('steps')
                result['cfg']       = inp.get('cfg')
                result['sampler']   = inp.get('sampler_name')
                result['scheduler'] = inp.get('scheduler')

            elif ct == 'CheckpointLoaderSimple':
                result['checkpoint'] = inp.get('ckpt_name')

            elif ct == 'CLIPTextEncode':
                # Identify positive vs negative by which KSampler input they connect to.
                # Simpler heuristic: node "6" = positive, node "7" = negative (our fixed workflow).
                # But to be robust, we check both via node graph references where possible.
                # Fallback: first CLIPTextEncode = positive, second = negative.
                text = inp.get('text', '')
                if result['positive'] is None:
                    result['positive'] = text
                elif result['negative'] is None:
                    result['negative'] = text

            elif ct == 'LoraLoader':
                name = inp.get('lora_name', '')
                if name:
                    result['loras'].append({
                        'name': name,
                        'strength': inp.get('strength_model', 1.0),
                    })

        # Re-parse with fixed node IDs from the known workflow (more reliable)
        try:
            if '6' in nodes:
                result['positive'] = nodes['6']['inputs'].get('text', result['positive'])
            if '7' in nodes:
                result['negative'] = nodes['7']['inputs'].get('text', result['negative'])
        except Exception:
            pass

        return result

    @app.route('/gallery')
    def gallery():
        category = request.args.get('category', 'all')
        slug = request.args.get('slug', '')
        sort = request.args.get('sort', 'newest')
        xref_category = request.args.get('xref_category', '')
        xref_slug = request.args.get('xref_slug', '')
        favourite_filter = request.args.get('favourite', '')
        nsfw_filter = request.args.get('nsfw', 'all')
        page = max(1, int(request.args.get('page', 1)))
        per_page = int(request.args.get('per_page', 48))
        per_page = per_page if per_page in (24, 48, 96) else 48

        images = _scan_gallery_images(category, slug)

        # Read sidecar data for filtering (favourite/NSFW/xref)
        upload_folder = app.config['UPLOAD_FOLDER']
        need_sidecar = (xref_category and xref_slug) or favourite_filter or nsfw_filter != 'all'
        if need_sidecar:
            for img in images:
                img['_sidecar'] = _read_sidecar(upload_folder, img['path']) or {}

        # Cross-reference filter
        if xref_category and xref_slug and xref_category in _XREF_KEY_MAP:
            sidecar_key = _XREF_KEY_MAP[xref_category]
            images = [img for img in images if img.get('_sidecar', {}).get(sidecar_key) == xref_slug]

        # Favourite filter
        if favourite_filter == 'on':
            images = [img for img in images if img.get('_sidecar', {}).get('is_favourite')]

        # NSFW filter
        if nsfw_filter == 'sfw':
            images = [img for img in images if not img.get('_sidecar', {}).get('is_nsfw')]
        elif nsfw_filter == 'nsfw':
            images = [img for img in images if img.get('_sidecar', {}).get('is_nsfw')]

        if sort == 'oldest':
            images.reverse()
        elif sort == 'random':
            import random
            random.shuffle(images)

        # Sort favourites first when favourite filter not active but sort is newest/oldest
        if sort in ('newest', 'oldest') and not favourite_filter and need_sidecar:
            images.sort(key=lambda x: (not x.get('_sidecar', {}).get('is_favourite', False), images.index(x)))

        total = len(images)
        total_pages = max(1, (total + per_page - 1) // per_page)
        page = min(page, total_pages)
        page_images = images[(page - 1) * per_page: page * per_page]
        _enrich_with_names(page_images)
        
        # Enrich with metadata for Info view
        upload_folder = os.path.abspath(app.config['UPLOAD_FOLDER'])
        for img in page_images:
            abs_img = os.path.join(upload_folder, img['path'])
            if os.path.isfile(abs_img) and abs_img.lower().endswith('.png'):
                img['meta'] = _parse_comfy_png_metadata(abs_img)
            else:
                img['meta'] = {}

        slug_options = []
        if category != 'all':
            Model = _MODEL_MAP.get(category)
            if Model:
                slug_options = [(r.slug, r.name) for r in Model.query.order_by(Model.name).with_entities(Model.slug, Model.name).all()]

        # Attach sidecar data to page images for template use
        for img in page_images:
            if '_sidecar' not in img:
                img['_sidecar'] = _read_sidecar(os.path.abspath(app.config['UPLOAD_FOLDER']), img['path']) or {}

        return render_template(
            'gallery.html',
            images=page_images,
            page=page,
            per_page=per_page,
            total=total,
            total_pages=total_pages,
            category=category,
            slug=slug,
            sort=sort,
            categories=GALLERY_CATEGORIES,
            slug_options=slug_options,
            xref_category=xref_category,
            xref_slug=xref_slug,
            favourite_filter=favourite_filter,
            nsfw_filter=nsfw_filter,
        )

    @app.route('/gallery/prompt-data')
    def gallery_prompt_data():
        """Return generation metadata for a specific image by reading its PNG tEXt chunk."""
        img_path = request.args.get('path', '')
        if not img_path:
            return {'error': 'path parameter required'}, 400

        # Validate path stays within uploads folder
        upload_folder = os.path.abspath(app.config['UPLOAD_FOLDER'])
        abs_img = os.path.abspath(os.path.join(upload_folder, img_path))
        if not abs_img.startswith(upload_folder + os.sep):
            return {'error': 'Invalid path'}, 400
        if not os.path.isfile(abs_img):
            return {'error': 'File not found'}, 404

        meta = _parse_comfy_png_metadata(abs_img)
        meta['path'] = img_path

        # Include sidecar data if available (for cross-reference links)
        sidecar = _read_sidecar(upload_folder, img_path)
        if sidecar:
            meta['sidecar'] = sidecar

        return meta

    def _write_sidecar(upload_folder, image_path, data):
        """Write/update JSON sidecar for an image."""
        sidecar = image_path.rsplit('.', 1)[0] + '.json'
        sidecar_path = os.path.join(upload_folder, sidecar)
        existing = {}
        try:
            with open(sidecar_path) as f:
                existing = json.load(f)
        except (OSError, json.JSONDecodeError):
            pass
        existing.update(data)
        with open(sidecar_path, 'w') as f:
            json.dump(existing, f, indent=2)

    @app.route('/gallery/image/favourite', methods=['POST'])
    def gallery_image_favourite():
        """Toggle favourite on a gallery image via sidecar JSON."""
        data = request.get_json(silent=True) or {}
        img_path = data.get('path', '')
        if not img_path:
            return {'error': 'path required'}, 400
        upload_folder = os.path.abspath(app.config['UPLOAD_FOLDER'])
        abs_img = os.path.abspath(os.path.join(upload_folder, img_path))
        if not abs_img.startswith(upload_folder + os.sep) or not os.path.isfile(abs_img):
            return {'error': 'Invalid path'}, 400
        sidecar = _read_sidecar(upload_folder, img_path) or {}
        new_val = not sidecar.get('is_favourite', False)
        _write_sidecar(upload_folder, img_path, {'is_favourite': new_val})
        return {'success': True, 'is_favourite': new_val}

    @app.route('/gallery/image/nsfw', methods=['POST'])
    def gallery_image_nsfw():
        """Toggle NSFW on a gallery image via sidecar JSON."""
        data = request.get_json(silent=True) or {}
        img_path = data.get('path', '')
        if not img_path:
            return {'error': 'path required'}, 400
        upload_folder = os.path.abspath(app.config['UPLOAD_FOLDER'])
        abs_img = os.path.abspath(os.path.join(upload_folder, img_path))
        if not abs_img.startswith(upload_folder + os.sep) or not os.path.isfile(abs_img):
            return {'error': 'Invalid path'}, 400
        sidecar = _read_sidecar(upload_folder, img_path) or {}
        new_val = not sidecar.get('is_nsfw', False)
        _write_sidecar(upload_folder, img_path, {'is_nsfw': new_val})
        return {'success': True, 'is_nsfw': new_val}

    @app.route('/gallery/delete', methods=['POST'])
    def gallery_delete():
        """Delete a generated image from the gallery. Only the image file is removed."""
        data = request.get_json(silent=True) or {}
        img_path = data.get('path', '')

        if not img_path:
            return {'error': 'path required'}, 400

        if len(img_path.split('/')) != 3:
            return {'error': 'invalid path format'}, 400

        upload_folder = os.path.abspath(app.config['UPLOAD_FOLDER'])
        abs_img = os.path.abspath(os.path.join(upload_folder, img_path))
        if not abs_img.startswith(upload_folder + os.sep):
            return {'error': 'Invalid path'}, 400

        if os.path.isfile(abs_img):
            os.remove(abs_img)
            # Also remove sidecar JSON if present
            sidecar = abs_img.rsplit('.', 1)[0] + '.json'
            if os.path.isfile(sidecar):
                os.remove(sidecar)

        return {'status': 'ok'}

    @app.route('/resource/<category>/<slug>/delete', methods=['POST'])
    def resource_delete(category, slug):
        """Delete a resource item from a category gallery.

        soft: removes JSON data file + DB record; LoRA/checkpoint file kept on disk.
        hard: removes JSON data file + LoRA/checkpoint safetensors + DB record.
        """
        _RESOURCE_MODEL_MAP = {
            'characters':  Character,
            'looks':       Look,
            'styles':      Style,
            'actions':     Action,
            'outfits':     Outfit,
            'scenes':      Scene,
            'detailers':   Detailer,
            'checkpoints': Checkpoint,
        }
        _RESOURCE_DATA_DIRS = {
            'characters':  app.config['CHARACTERS_DIR'],
            'looks':       app.config['LOOKS_DIR'],
            'styles':      app.config['STYLES_DIR'],
            'actions':     app.config['ACTIONS_DIR'],
            'outfits':     app.config['CLOTHING_DIR'],
            'scenes':      app.config['SCENES_DIR'],
            'detailers':   app.config['DETAILERS_DIR'],
            'checkpoints': app.config['CHECKPOINTS_DIR'],
        }
        _LORA_BASE = '/ImageModels/lora/'

        if category not in _RESOURCE_MODEL_MAP:
            return {'error': 'unknown category'}, 400

        req = request.get_json(silent=True) or {}
        mode = req.get('mode', 'soft')

        data_dir = _RESOURCE_DATA_DIRS[category]
        json_path = os.path.join(data_dir, f'{slug}.json')

        deleted = []
        asset_abs = None

        # Resolve asset path before deleting JSON (hard only)
        if mode == 'hard' and os.path.isfile(json_path):
            try:
                with open(json_path) as f:
                    item_data = json.load(f)
                if category == 'checkpoints':
                    ckpt_rel = item_data.get('checkpoint_path', '')
                    if ckpt_rel.startswith('Illustrious/'):
                        asset_abs = os.path.join(app.config['ILLUSTRIOUS_MODELS_DIR'],
                                                 ckpt_rel[len('Illustrious/'):])
                    elif ckpt_rel.startswith('Noob/'):
                        asset_abs = os.path.join(app.config['NOOB_MODELS_DIR'],
                                                 ckpt_rel[len('Noob/'):])
                else:
                    lora_name = item_data.get('lora', {}).get('lora_name', '')
                    if lora_name:
                        asset_abs = os.path.join(_LORA_BASE, lora_name)
            except Exception:
                pass

        # Delete JSON
        if os.path.isfile(json_path):
            os.remove(json_path)
            deleted.append('json')

        # Delete LoRA/checkpoint file (hard only)
        if mode == 'hard' and asset_abs and os.path.isfile(asset_abs):
            os.remove(asset_abs)
            deleted.append('lora' if category != 'checkpoints' else 'checkpoint')

        # Remove DB record
        Model = _RESOURCE_MODEL_MAP[category]
        rec = Model.query.filter_by(slug=slug).first()
        if rec:
            db.session.delete(rec)
            db.session.commit()
            deleted.append('db')

        return {'status': 'ok', 'deleted': deleted}
