import json
import logging
from flask import render_template, request, session
from models import Character, Action, Style, Scene, Detailer, Checkpoint
from services.prompts import build_multi_prompt, build_extras_prompt, _resolve_character
from services.workflow import _prepare_workflow
from services.job_queue import _enqueue_job, _make_finalize
from services.file_io import get_available_checkpoints

logger = logging.getLogger('gaze')


def register_routes(app):

    @app.route('/multi', methods=['GET', 'POST'])
    def multi_char():
        characters = Character.query.order_by(Character.name).all()
        checkpoints = get_available_checkpoints()
        actions = Action.query.order_by(Action.name).all()
        scenes = Scene.query.order_by(Scene.name).all()
        styles = Style.query.order_by(Style.name).all()
        detailers = Detailer.query.order_by(Detailer.name).all()

        if not checkpoints:
            checkpoints = ["Noob/oneObsession_v19Atypical.safetensors"]

        if request.method == 'POST':
            char_a_slug = request.form.get('char_a')
            char_b_slug = request.form.get('char_b')
            checkpoint = request.form.get('checkpoint')
            custom_positive = request.form.get('positive_prompt', '')
            custom_negative = request.form.get('negative_prompt', '')

            action_slugs = request.form.getlist('action_slugs')
            scene_slugs = request.form.getlist('scene_slugs')
            style_slugs = request.form.getlist('style_slugs')
            detailer_slugs = request.form.getlist('detailer_slugs')
            width = request.form.get('width') or 1024
            height = request.form.get('height') or 1024

            char_a = _resolve_character(char_a_slug)
            char_b = _resolve_character(char_b_slug)

            if not char_a or not char_b:
                if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                    return {'error': 'Both characters must be selected'}, 400
                return render_template('multi_char.html', characters=characters, checkpoints=checkpoints,
                                       actions=actions, scenes=scenes, styles=styles, detailers=detailers)

            sel_actions = Action.query.filter(Action.slug.in_(action_slugs)).all() if action_slugs else []
            sel_scenes = Scene.query.filter(Scene.slug.in_(scene_slugs)).all() if scene_slugs else []
            sel_styles = Style.query.filter(Style.slug.in_(style_slugs)).all() if style_slugs else []
            sel_detailers = Detailer.query.filter(Detailer.slug.in_(detailer_slugs)).all() if detailer_slugs else []

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

                # Build extras from mix-and-match selections
                extras = build_extras_prompt(sel_actions, [], sel_scenes, sel_styles, sel_detailers)
                if custom_positive:
                    extras = f"{custom_positive}, {extras}" if extras else custom_positive

                # Build multi-character prompt
                prompts = build_multi_prompt(char_a, char_b, extras_prompt=extras)

                # Apply prompt overrides if provided
                override_main = request.form.get('override_prompt', '').strip()
                if override_main:
                    prompts['main'] = override_main
                for key in ('char_a_main', 'char_b_main', 'char_a_face', 'char_b_face'):
                    override = request.form.get(f'override_{key}', '').strip()
                    if override:
                        prompts[key] = override

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

                # Prepare workflow with both character LoRAs
                ckpt_obj = Checkpoint.query.filter_by(checkpoint_path=checkpoint).first() if checkpoint else None
                workflow = _prepare_workflow(
                    workflow, char_a, prompts, checkpoint, custom_negative,
                    action=sel_actions[0] if sel_actions else None,
                    style=sel_styles[0] if sel_styles else None,
                    detailer=sel_detailers[0] if sel_detailers else None,
                    scene=sel_scenes[0] if sel_scenes else None,
                    width=width,
                    height=height,
                    checkpoint_data=ckpt_obj.data if ckpt_obj else None,
                    fixed_seed=fixed_seed,
                    character_b=char_b,
                )

                label = f"Multi: {char_a.name} + {char_b.name}"
                slug_pair = f"{char_a.slug}_{char_b.slug}"
                _finalize = _make_finalize('multi', slug_pair)
                job = _enqueue_job(label, workflow, _finalize)

                # Save to session
                session['multi_char_a'] = char_a.slug
                session['multi_char_b'] = char_b.slug
                session.modified = True

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

            except Exception as e:
                logger.exception("Multi-char generation error: %s", e)
                if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                    return {'error': str(e)}, 500

        return render_template('multi_char.html', characters=characters, checkpoints=checkpoints,
                               actions=actions, scenes=scenes, styles=styles, detailers=detailers,
                               selected_char_a=session.get('multi_char_a', ''),
                               selected_char_b=session.get('multi_char_b', ''))

    @app.route('/multi/preview_prompt', methods=['POST'])
    def multi_preview_prompt():
        char_a_slug = request.form.get('char_a')
        char_b_slug = request.form.get('char_b')
        if not char_a_slug or not char_b_slug:
            return {'error': 'Both characters must be selected'}, 400

        char_a = _resolve_character(char_a_slug)
        char_b = _resolve_character(char_b_slug)
        if not char_a or not char_b:
            return {'error': 'Character not found'}, 404

        action_slugs = request.form.getlist('action_slugs')
        scene_slugs = request.form.getlist('scene_slugs')
        style_slugs = request.form.getlist('style_slugs')
        detailer_slugs = request.form.getlist('detailer_slugs')
        custom_positive = request.form.get('positive_prompt', '')

        sel_actions = Action.query.filter(Action.slug.in_(action_slugs)).all() if action_slugs else []
        sel_scenes = Scene.query.filter(Scene.slug.in_(scene_slugs)).all() if scene_slugs else []
        sel_styles = Style.query.filter(Style.slug.in_(style_slugs)).all() if style_slugs else []
        sel_detailers = Detailer.query.filter(Detailer.slug.in_(detailer_slugs)).all() if detailer_slugs else []

        extras = build_extras_prompt(sel_actions, [], sel_scenes, sel_styles, sel_detailers)
        if custom_positive:
            extras = f"{custom_positive}, {extras}" if extras else custom_positive

        prompts = build_multi_prompt(char_a, char_b, extras_prompt=extras)
        return {
            'prompt': prompts['main'],
            'hand': prompts['hand'],
            'char_a_main': prompts['char_a_main'],
            'char_a_face': prompts['char_a_face'],
            'char_b_main': prompts['char_b_main'],
            'char_b_face': prompts['char_b_face'],
        }
