import json
import os
import re
import logging

from flask import render_template, request, redirect, url_for, flash, session
from sqlalchemy.orm.attributes import flag_modified

from models import db, Character, Scene, Settings
from services.workflow import _prepare_workflow, _get_default_checkpoint
from services.job_queue import _enqueue_job, _make_finalize, _enqueue_task
from services.prompts import build_prompt, _resolve_character, _ensure_character_fields
from services.sync import sync_scenes
from services.file_io import get_available_loras
from services.llm import load_prompt, call_llm
from routes.shared import register_common_routes
from utils import _WARDROBE_KEYS

logger = logging.getLogger('gaze')


def register_routes(app):
    register_common_routes(app, 'scenes')

    @app.route('/scenes')
    def scenes_index():
        query = Scene.query
        fav = request.args.get('favourite')
        nsfw = request.args.get('nsfw', 'all')
        if fav == 'on':
            query = query.filter_by(is_favourite=True)
        if nsfw == 'sfw':
            query = query.filter_by(is_nsfw=False)
        elif nsfw == 'nsfw':
            query = query.filter_by(is_nsfw=True)
        scenes = query.order_by(Scene.is_favourite.desc(), Scene.name).all()
        return render_template('scenes/index.html', scenes=scenes, favourite_filter=fav or '', nsfw_filter=nsfw)

    @app.route('/scenes/rescan', methods=['POST'])
    def rescan_scenes():
        sync_scenes()
        flash('Database synced with scene files.')
        return redirect(url_for('scenes_index'))

    @app.route('/scene/<path:slug>')
    def scene_detail(slug):
        scene = Scene.query.filter_by(slug=slug).first_or_404()
        characters = Character.query.order_by(Character.name).all()

        # Load state from session
        preferences = session.get(f'prefs_scene_{slug}')
        preview_image = session.get(f'preview_scene_{slug}')
        selected_character = session.get(f'char_scene_{slug}')
        extra_positive = session.get(f'extra_pos_scene_{slug}', '')
        extra_negative = session.get(f'extra_neg_scene_{slug}', '')

        # List existing preview images
        upload_dir = os.path.join(app.config['UPLOAD_FOLDER'], f"scenes/{slug}")
        existing_previews = []
        if os.path.isdir(upload_dir):
            files = sorted([f for f in os.listdir(upload_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))], reverse=True)
            existing_previews = [f"scenes/{slug}/{f}" for f in files]

        return render_template('scenes/detail.html', scene=scene, characters=characters,
                              preferences=preferences, preview_image=preview_image,
                              selected_character=selected_character, existing_previews=existing_previews,
                              extra_positive=extra_positive, extra_negative=extra_negative)

    @app.route('/scene/<path:slug>/edit', methods=['GET', 'POST'])
    def edit_scene(slug):
        scene = Scene.query.filter_by(slug=slug).first_or_404()
        loras = get_available_loras('scenes')

        if request.method == 'POST':
            try:
                # 1. Update basic fields
                scene.name = request.form.get('scene_name')

                # 2. Rebuild the data dictionary
                new_data = scene.data.copy()
                new_data['scene_name'] = scene.name

                # Update scene section
                if 'scene' in new_data:
                    for key in new_data['scene'].keys():
                        form_key = f"scene_{key}"
                        if form_key in request.form:
                            val = request.form.get(form_key)
                            # Handle list for furniture/colors if they were originally lists
                            if key in ['furniture', 'colors'] and isinstance(new_data['scene'][key], list):
                                val = [v.strip() for v in val.split(',') if v.strip()]
                            new_data['scene'][key] = val

                # Update lora section
                if 'lora' in new_data:
                    for key in new_data['lora'].keys():
                        form_key = f"lora_{key}"
                        if form_key in request.form:
                            val = request.form.get(form_key)
                            if key == 'lora_weight':
                                try: val = float(val)
                                except: val = 1.0
                            new_data['lora'][key] = val

                # LoRA weight randomization bounds
                for bound in ['lora_weight_min', 'lora_weight_max']:
                    val_str = request.form.get(f'lora_{bound}', '').strip()
                    if val_str:
                        try:
                            new_data.setdefault('lora', {})[bound] = float(val_str)
                        except ValueError:
                            pass
                    else:
                        new_data.setdefault('lora', {}).pop(bound, None)

                # Update Tags (structured dict)
                new_data['tags'] = {
                    'scene_type': request.form.get('tag_scene_type', '').strip(),
                    'nsfw': 'tag_nsfw' in request.form,
                }
                scene.is_nsfw = new_data['tags']['nsfw']

                scene.data = new_data
                flag_modified(scene, "data")

                # 3. Write back to JSON file
                scene_file = scene.filename or f"{re.sub(r'[^a-zA-Z0-9_]', '', scene.scene_id)}.json"
                file_path = os.path.join(app.config['SCENES_DIR'], scene_file)

                with open(file_path, 'w') as f:
                    json.dump(new_data, f, indent=2)

                db.session.commit()
                flash('Scene updated successfully!')
                return redirect(url_for('scene_detail', slug=slug))

            except Exception as e:
                logger.exception("Edit error: %s", e)
                flash(f"Error saving changes: {str(e)}")

        return render_template('scenes/edit.html', scene=scene, loras=loras)

    def _queue_scene_generation(scene_obj, character=None, selected_fields=None, client_id=None, fixed_seed=None, extra_positive=None, extra_negative=None):
        if character:
            combined_data = character.data.copy()
            combined_data['character_id'] = character.character_id

            # Update character's 'defaults' with scene details
            scene_data = scene_obj.data.get('scene', {})

            # Build scene tag string
            scene_tags = []
            for key in ['background', 'foreground', 'furniture', 'colors', 'lighting', 'theme']:
                val = scene_data.get(key)
                if val:
                    if isinstance(val, list):
                        scene_tags.extend(val)
                    else:
                        scene_tags.append(val)

            combined_data['defaults']['scene'] = ", ".join(scene_tags)

            # Merge scene lora triggers if present
            scene_lora = scene_obj.data.get('lora', {})
            if scene_lora.get('lora_triggers'):
                if 'lora' not in combined_data: combined_data['lora'] = {}
                combined_data['lora']['lora_triggers'] = f"{combined_data['lora'].get('lora_triggers', '')}, {scene_lora['lora_triggers']}"

            # Merge character identity and wardrobe fields into selected_fields
            if selected_fields:
                _ensure_character_fields(character, selected_fields)
            else:
                # Auto-include essential character fields (minimal set for batch/default generation)
                selected_fields = []
                for key in ['base', 'head']:
                    if character.data.get('identity', {}).get(key):
                        selected_fields.append(f'identity::{key}')
                selected_fields.append('special::name')
                wardrobe = character.get_active_wardrobe()
                for key in _WARDROBE_KEYS:
                    if wardrobe.get(key):
                        selected_fields.append(f'wardrobe::{key}')
                selected_fields.extend(['defaults::scene', 'lora::lora_triggers'])

            default_fields = scene_obj.default_fields
            active_outfit = character.active_outfit
        else:
            # Scene only - no character
            scene_data = scene_obj.data.get('scene', {})
            scene_tags = []
            for key in ['background', 'foreground', 'furniture', 'colors', 'lighting', 'theme']:
                val = scene_data.get(key)
                if val:
                    if isinstance(val, list): scene_tags.extend(val)
                    else: scene_tags.append(val)

            combined_data = {
                'character_id': scene_obj.scene_id,
                'defaults': {
                    'scene': ", ".join(scene_tags)
                },
                'lora': scene_obj.data.get('lora', {}),
                'tags': scene_obj.data.get('tags', [])
            }
            if not selected_fields:
                selected_fields = ['defaults::scene', 'lora::lora_triggers']
            default_fields = scene_obj.default_fields
            active_outfit = 'default'

        with open('comfy_workflow.json', 'r') as f:
            workflow = json.load(f)

        prompts = build_prompt(combined_data, selected_fields, default_fields, active_outfit=active_outfit)

        if extra_positive:
            prompts["main"] = f"{prompts['main']}, {extra_positive}"

        # For scene generation, we want to ensure Node 20 is handled in _prepare_workflow
        ckpt_path, ckpt_data = _get_default_checkpoint()
        workflow = _prepare_workflow(workflow, character, prompts, scene=scene_obj, custom_negative=extra_negative or None, checkpoint=ckpt_path, checkpoint_data=ckpt_data, fixed_seed=fixed_seed)
        return workflow

    @app.route('/scene/<path:slug>/generate', methods=['POST'])
    def generate_scene_image(slug):
        scene_obj = Scene.query.filter_by(slug=slug).first_or_404()

        try:
            # Get action type
            action = request.form.get('action', 'preview')

            # Get selected fields
            selected_fields = request.form.getlist('include_field')

            # Get selected character (if any)
            character_slug = request.form.get('character_slug', '')
            character = _resolve_character(character_slug)
            if character_slug == '__random__' and character:
                character_slug = character.slug

            # Get additional prompts
            extra_positive = request.form.get('extra_positive', '').strip()
            extra_negative = request.form.get('extra_negative', '').strip()

            # Save preferences
            session[f'char_scene_{slug}'] = character_slug
            session[f'prefs_scene_{slug}'] = selected_fields
            session[f'extra_pos_scene_{slug}'] = extra_positive
            session[f'extra_neg_scene_{slug}'] = extra_negative
            session.modified = True

            # Parse optional seed
            seed_val = request.form.get('seed', '').strip()
            fixed_seed = int(seed_val) if seed_val else None

            # Build workflow using helper
            workflow = _queue_scene_generation(scene_obj, character, selected_fields, fixed_seed=fixed_seed, extra_positive=extra_positive, extra_negative=extra_negative)

            char_label = character.name if character else 'no character'
            label = f"Scene: {scene_obj.name} ({char_label}) – {action}"
            job = _enqueue_job(label, workflow, _make_finalize('scenes', slug, Scene, action))

            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return {'status': 'queued', 'job_id': job['id']}

            return redirect(url_for('scene_detail', slug=slug))

        except Exception as e:
            logger.exception("Generation error: %s", e)
            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return {'error': str(e)}, 500
            flash(f"Error during generation: {str(e)}")
            return redirect(url_for('scene_detail', slug=slug))

    @app.route('/scenes/bulk_create', methods=['POST'])
    def bulk_create_scenes_from_loras():
        _s = Settings.query.first()
        backgrounds_lora_dir = ((_s.lora_dir_scenes if _s else None) or '/ImageModels/lora/Illustrious/Backgrounds').rstrip('/')
        _lora_subfolder = os.path.basename(backgrounds_lora_dir)
        if not os.path.exists(backgrounds_lora_dir):
            flash('Backgrounds LoRA directory not found.', 'error')
            return redirect(url_for('scenes_index'))

        overwrite = request.form.get('overwrite') == 'true'
        skipped = 0
        job_ids = []

        system_prompt = load_prompt('scene_system.txt')
        if not system_prompt:
            flash('Scene system prompt file not found.', 'error')
            return redirect(url_for('scenes_index'))

        scenes_dir = app.config['SCENES_DIR']

        for filename in os.listdir(backgrounds_lora_dir):
            if filename.endswith('.safetensors'):
                name_base = filename.rsplit('.', 1)[0]
                scene_id = re.sub(r'[^a-zA-Z0-9_]', '_', name_base.lower())
                scene_name = re.sub(r'[^a-zA-Z0-9]+', ' ', name_base).title()

                json_filename = f"{scene_id}.json"
                json_path = os.path.join(scenes_dir, json_filename)

                is_existing = os.path.exists(json_path)
                if is_existing and not overwrite:
                    skipped += 1
                    continue

                html_filename = f"{name_base}.html"
                html_path = os.path.join(backgrounds_lora_dir, html_filename)
                html_content = ""
                if os.path.exists(html_path):
                    try:
                        with open(html_path, 'r', encoding='utf-8', errors='ignore') as hf:
                            html_raw = hf.read()
                            clean_html = re.sub(r'<script[^>]*>.*?</script>', '', html_raw, flags=re.DOTALL)
                            clean_html = re.sub(r'<style[^>]*>.*?</style>', '', clean_html, flags=re.DOTALL)
                            clean_html = re.sub(r'<img[^>]*>', '', clean_html)
                            clean_html = re.sub(r'<[^>]+>', ' ', clean_html)
                            html_content = ' '.join(clean_html.split())
                    except Exception as e:
                        logger.error("Error reading HTML %s: %s", html_filename, e)

                def make_task(filename, name_base, scene_id, scene_name, json_path, _lora_subfolder, html_content, system_prompt, is_existing):
                    def task_fn(job):
                        prompt = f"Describe a scene for an AI image generation model based on the LoRA filename: '{filename}'"
                        if html_content:
                            prompt += f"\n\nHere is descriptive text and metadata extracted from an associated HTML file for this LoRA:\n###\n{html_content[:3000]}\n###"

                        llm_response = call_llm(prompt, system_prompt)
                        clean_json = llm_response.replace('```json', '').replace('```', '').strip()
                        scene_data = json.loads(clean_json)

                        scene_data['scene_id'] = scene_id
                        scene_data['scene_name'] = scene_name

                        if 'lora' not in scene_data: scene_data['lora'] = {}
                        scene_data['lora']['lora_name'] = f"Illustrious/{_lora_subfolder}/{filename}"

                        if not scene_data['lora'].get('lora_triggers'):
                            scene_data['lora']['lora_triggers'] = name_base
                        if scene_data['lora'].get('lora_weight') is None:
                            scene_data['lora']['lora_weight'] = 1.0
                        if scene_data['lora'].get('lora_weight_min') is None:
                            scene_data['lora']['lora_weight_min'] = 0.7
                        if scene_data['lora'].get('lora_weight_max') is None:
                            scene_data['lora']['lora_weight_max'] = 1.0

                        with open(json_path, 'w') as f:
                            json.dump(scene_data, f, indent=2)

                        job['result'] = {'name': scene_name, 'action': 'overwritten' if is_existing else 'created'}
                    return task_fn

                job = _enqueue_task(f"Create scene: {scene_name}", make_task(filename, name_base, scene_id, scene_name, json_path, _lora_subfolder, html_content, system_prompt, is_existing))
                job_ids.append(job['id'])

        if job_ids:
            def sync_task(job):
                sync_scenes()
                job['result'] = {'synced': True}
            _enqueue_task("Sync scenes DB", sync_task)

        if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
            return {'success': True, 'queued': len(job_ids), 'skipped': skipped}
        flash(f'Queued {len(job_ids)} scene tasks ({skipped} skipped).')
        return redirect(url_for('scenes_index'))

    @app.route('/scene/create', methods=['GET', 'POST'])
    def create_scene():
        form_data = {}

        if request.method == 'POST':
            name = request.form.get('name')
            slug = request.form.get('filename', '').strip()

            form_data = {'name': name, 'filename': slug}

            if not slug:
                slug = re.sub(r'[^a-zA-Z0-9]+', '_', name.lower()).strip('_')

            safe_slug = re.sub(r'[^a-zA-Z0-9_]', '', slug)
            if not safe_slug:
                safe_slug = 'scene'

            base_slug = safe_slug
            counter = 1
            while os.path.exists(os.path.join(app.config['SCENES_DIR'], f"{safe_slug}.json")):
                safe_slug = f"{base_slug}_{counter}"
                counter += 1

            scene_data = {
                "scene_id": safe_slug,
                "scene_name": name,
                "scene": {
                    "background": "",
                    "foreground": "",
                    "furniture": [],
                    "colors": [],
                    "lighting": "",
                    "theme": ""
                },
                "lora": {
                    "lora_name": "",
                    "lora_weight": 1.0,
                    "lora_triggers": ""
                }
            }

            try:
                file_path = os.path.join(app.config['SCENES_DIR'], f"{safe_slug}.json")
                with open(file_path, 'w') as f:
                    json.dump(scene_data, f, indent=2)

                new_scene = Scene(
                    scene_id=safe_slug, slug=safe_slug, filename=f"{safe_slug}.json",
                    name=name, data=scene_data
                )
                db.session.add(new_scene)
                db.session.commit()

                flash('Scene created successfully!')
                return redirect(url_for('scene_detail', slug=safe_slug))
            except Exception as e:
                logger.exception("Save error: %s", e)
                flash(f"Failed to create scene: {e}")
                return render_template('scenes/create.html', form_data=form_data)

        return render_template('scenes/create.html', form_data=form_data)

