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, Look, 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, _dedup_tags
from services.sync import sync_looks
from services.file_io import get_available_loras, _count_look_assignments
from services.llm import load_prompt, call_llm
from routes.shared import register_common_routes

logger = logging.getLogger('gaze')


def _ensure_look_lora_prefix(lora_name):
    """Ensure look LoRA paths have the correct 'Illustrious/Looks/' prefix."""
    if not lora_name:
        return lora_name
    
    if not lora_name.startswith('Illustrious/Looks/'):
        # Add the prefix if missing
        if lora_name.startswith('Illustrious/'):
            # Has Illustrious but wrong subfolder - replace
            parts = lora_name.split('/', 1)
            if len(parts) > 1:
                lora_name = 'Illustrious/Looks/' + parts[1]
            else:
                lora_name = 'Illustrious/Looks/' + lora_name
        else:
            # No prefix at all - add it
            lora_name = 'Illustrious/Looks/' + lora_name
    
    return lora_name


def _fix_look_lora_data(lora_data):
    """Fix look LoRA data to ensure correct prefix."""
    if not lora_data:
        return lora_data
    
    lora_name = lora_data.get('lora_name', '')
    if lora_name:
        lora_data = lora_data.copy()  # Avoid mutating original
        lora_data['lora_name'] = _ensure_look_lora_prefix(lora_name)
    
    return lora_data


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

    @app.route('/looks')
    def looks_index():
        query = Look.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)
        looks = query.order_by(Look.is_favourite.desc(), Look.name).all()
        look_assignments = _count_look_assignments()
        return render_template('looks/index.html', looks=looks, look_assignments=look_assignments, favourite_filter=fav or '', nsfw_filter=nsfw)

    @app.route('/looks/rescan', methods=['POST'])
    def rescan_looks():
        sync_looks()
        flash('Database synced with look files.')
        return redirect(url_for('looks_index'))

    @app.route('/look/<path:slug>')
    def look_detail(slug):
        look = Look.query.filter_by(slug=slug).first_or_404()
        characters = Character.query.order_by(Character.name).all()

        # Pre-select the linked characters if set (supports multi-character assignment)
        preferences = session.get(f'prefs_look_{slug}')
        preview_image = session.get(f'preview_look_{slug}')

        # Get linked character IDs (new character_ids JSON field)
        linked_character_ids = look.character_ids or []
        # Fallback to legacy character_id if character_ids is empty
        if not linked_character_ids and look.character_id:
            linked_character_ids = [look.character_id]

        # Session-selected character for preview (single selection for generation)
        selected_character = session.get(f'char_look_{slug}', linked_character_ids[0] if linked_character_ids else '')

        # FIX: Add existing_previews scanning (matching other resource routes)
        upload_folder = app.config['UPLOAD_FOLDER']
        preview_dir = os.path.join(upload_folder, 'looks', slug)
        existing_previews = []
        if os.path.isdir(preview_dir):
            for f in os.listdir(preview_dir):
                if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
                    existing_previews.append(f'looks/{slug}/{f}')
        existing_previews.sort()

        extra_positive = session.get(f'extra_pos_look_{slug}', '')
        extra_negative = session.get(f'extra_neg_look_{slug}', '')

        return render_template('looks/detail.html', look=look, characters=characters,
                               preferences=preferences, preview_image=preview_image,
                               selected_character=selected_character,
                               linked_character_ids=linked_character_ids,
                               existing_previews=existing_previews,
                               extra_positive=extra_positive, extra_negative=extra_negative)

    @app.route('/look/<path:slug>/edit', methods=['GET', 'POST'])
    def edit_look(slug):
        look = Look.query.filter_by(slug=slug).first_or_404()
        characters = Character.query.order_by(Character.name).all()
        loras = get_available_loras('characters')

        if request.method == 'POST':
            look.name = request.form.get('look_name', look.name)

            # Handle multiple character IDs from checkboxes
            character_ids = request.form.getlist('character_ids')
            look.character_ids = character_ids if character_ids else []

            # Also update legacy character_id field for backward compatibility
            if character_ids:
                look.character_id = character_ids[0]
            else:
                look.character_id = None

            new_data = look.data.copy()
            new_data['look_name'] = look.name
            new_data['character_id'] = look.character_id

            new_data['positive'] = request.form.get('positive', '')
            new_data['negative'] = request.form.get('negative', '')

            lora_name = request.form.get('lora_lora_name', '')
            lora_weight = float(request.form.get('lora_lora_weight', 1.0) or 1.0)
            lora_triggers = request.form.get('lora_lora_triggers', '')
            new_data['lora'] = {'lora_name': lora_name, 'lora_weight': lora_weight, 'lora_triggers': lora_triggers}
            for bound in ['lora_weight_min', 'lora_weight_max']:
                val_str = request.form.get(f'lora_{bound}', '').strip()
                if val_str:
                    try:
                        new_data['lora'][bound] = float(val_str)
                    except ValueError:
                        pass

            new_data['tags'] = {
                'origin_series': request.form.get('tag_origin_series', '').strip(),
                'origin_type': request.form.get('tag_origin_type', '').strip(),
                'nsfw': 'tag_nsfw' in request.form,
            }
            look.is_nsfw = new_data['tags']['nsfw']

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

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

            flash(f'Look "{look.name}" updated!')
            return redirect(url_for('look_detail', slug=look.slug))

        return render_template('looks/edit.html', look=look, characters=characters, loras=loras)

    @app.route('/look/<path:slug>/generate', methods=['POST'])
    def generate_look_image(slug):
        look = Look.query.filter_by(slug=slug).first_or_404()

        try:
            action = request.form.get('action', 'preview')
            selected_fields = request.form.getlist('include_field')

            character_slug = request.form.get('character_slug', '')
            character = None

            # Only load a character when the user explicitly selects one
            character = _resolve_character(character_slug)
            if character_slug == '__random__' and character:
                character_slug = character.slug
            elif character_slug and not character:
                # fallback: try matching by character_id
                character = Character.query.filter_by(character_id=character_slug).first()
            # No fallback to look.character_id — looks are self-contained

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

            session[f'prefs_look_{slug}'] = selected_fields
            session[f'char_look_{slug}'] = character_slug
            session[f'extra_pos_look_{slug}'] = extra_positive
            session[f'extra_neg_look_{slug}'] = extra_negative
            session.modified = True

            lora_triggers = look.data.get('lora', {}).get('lora_triggers', '')
            look_positive = look.data.get('positive', '')

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

            if character:
                # Merge character identity with look LoRA and positive prompt
                combined_data = {
                    'character_id': character.character_id,
                    'identity': character.data.get('identity', {}),
                    'defaults': character.data.get('defaults', {}),
                    'wardrobe': character.data.get('wardrobe', {}).get(character.active_outfit or 'default',
                                character.data.get('wardrobe', {}).get('default', {})),
                    'styles': character.data.get('styles', {}),
                    'lora': _fix_look_lora_data(look.data.get('lora', {})),
                    'tags': look.data.get('tags', [])
                }
                _ensure_character_fields(character, selected_fields,
                                         include_wardrobe=False, include_defaults=True)
                prompts = build_prompt(combined_data, selected_fields, character.default_fields)
                # Append look-specific triggers and positive
                extra = ', '.join(filter(None, [lora_triggers, look_positive]))
                if extra:
                    prompts['main'] = _dedup_tags(f"{prompts['main']}, {extra}" if prompts['main'] else extra)
                primary_color = character.data.get('styles', {}).get('primary_color', '')
                bg = f"{primary_color} simple background" if primary_color else "simple background"
            else:
                # Look is self-contained: build prompt from its own positive and triggers only
                main = _dedup_tags(', '.join(filter(None, ['(solo:1.2)', lora_triggers, look_positive])))
                prompts = {'main': main, 'face': '', 'hand': ''}
                bg = "simple background"

            prompts['main'] = _dedup_tags(f"{prompts['main']}, {bg}" if prompts['main'] else bg)

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

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

            ckpt_path, ckpt_data = _get_default_checkpoint()
            workflow = _prepare_workflow(workflow, character, prompts, custom_negative=extra_negative or None, checkpoint=ckpt_path,
                                         checkpoint_data=ckpt_data, look=look, fixed_seed=fixed_seed)

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

            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return {'status': 'queued', 'job_id': job['id']}
            return redirect(url_for('look_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('look_detail', slug=slug))

    @app.route('/look/<path:slug>/generate_character', methods=['POST'])
    def generate_character_from_look(slug):
        """Generate a character JSON using a look as the base."""
        look = Look.query.filter_by(slug=slug).first_or_404()

        # Get or validate inputs
        character_name = request.form.get('character_name', look.name)
        use_llm = request.form.get('use_llm') == 'on'

        # Auto-generate slug
        character_slug = re.sub(r'[^a-zA-Z0-9]+', '_', character_name.lower()).strip('_')
        character_slug = re.sub(r'[^a-zA-Z0-9_]', '', character_slug)

        # Find available filename
        base_slug = character_slug
        counter = 1
        while os.path.exists(os.path.join(app.config['CHARACTERS_DIR'], f"{character_slug}.json")):
            character_slug = f"{base_slug}_{counter}"
            counter += 1

        if use_llm:
            # Use LLM to generate character from look context
            system_prompt = load_prompt('character_system.txt')
            if not system_prompt:
                flash('Character system prompt file not found.', 'error')
                return redirect(url_for('look_detail', slug=slug))

            prompt = f"""Generate a character based on this look description:

Look Name: {look.name}
Positive Prompt: {look.data.get('positive', '')}
Negative Prompt: {look.data.get('negative', '')}
Tags: {', '.join(look.data.get('tags', []))}
LoRA Triggers: {look.data.get('lora', {}).get('lora_triggers', '')}

Create a complete character JSON with identity, styles, and appropriate wardrobe fields.
The character should match the visual style described in the look.

Character Name: {character_name}
Character ID: {character_slug}"""

            try:
                llm_response = call_llm(prompt, system_prompt)
                # Clean response (remove markdown if present)
                clean_json = llm_response.replace('```json', '').replace('```', '').strip()
                character_data = json.loads(clean_json)

                # Enforce IDs
                character_data['character_id'] = character_slug
                character_data['character_name'] = character_name
                
                # Ensure the character inherits the look's LoRA with correct path
                lora_data = _fix_look_lora_data(look.data.get('lora', {}).copy())
                character_data['lora'] = lora_data

            except Exception as e:
                logger.exception(f"LLM character generation error: {e}")
                flash(f'Failed to generate character with AI: {e}', 'error')
                return redirect(url_for('look_detail', slug=slug))
        else:
            # Create minimal character template
            lora_data = _fix_look_lora_data(look.data.get('lora', {}).copy())
            
            character_data = {
                "character_id": character_slug,
                "character_name": character_name,
                "identity": {
                    "base": lora_data.get('lora_triggers', ''),
                    "head": "",
                    "upper_body": "",
                    "lower_body": "",
                    "hands": "",
                    "feet": "",
                    "additional": ""
                },
                "defaults": {
                    "expression": "",
                    "pose": "",
                    "scene": ""
                },
                "wardrobe": {
                    "base": "",
                    "head": "",
                    "upper_body": "",
                    "lower_body": "",
                    "hands": "",
                    "feet": "",
                    "additional": ""
                },
                "styles": {
                    "aesthetic": "",
                    "primary_color": "",
                    "secondary_color": "",
                    "tertiary_color": ""
                },
                "lora": lora_data,
                "tags": look.data.get('tags', [])
            }

        # Save character JSON
        char_path = os.path.join(app.config['CHARACTERS_DIR'], f"{character_slug}.json")
        try:
            with open(char_path, 'w') as f:
                json.dump(character_data, f, indent=2)
        except Exception as e:
            flash(f'Failed to save character file: {e}', 'error')
            return redirect(url_for('look_detail', slug=slug))

        # Create DB entry
        character = Character(
            character_id=character_slug,
            slug=character_slug,
            name=character_name,
            data=character_data
        )
        db.session.add(character)
        db.session.commit()

        # Link the look to this character
        look.character_id = character_slug
        db.session.commit()

        flash(f'Character "{character_name}" created from look!', 'success')
        return redirect(url_for('detail', slug=character_slug))

    @app.route('/look/create', methods=['GET', 'POST'])
    def create_look():
        characters = Character.query.order_by(Character.name).all()
        loras = get_available_loras('characters')
        form_data = {}

        if request.method == 'POST':
            name = request.form.get('name', '').strip()
            character_id = request.form.get('character_id', '') or None
            lora_name = request.form.get('lora_lora_name', '')
            lora_weight = float(request.form.get('lora_lora_weight', 1.0) or 1.0)
            lora_triggers = request.form.get('lora_lora_triggers', '')
            positive = request.form.get('positive', '')
            negative = request.form.get('negative', '')
            tags = {
                'origin_series': request.form.get('tag_origin_series', '').strip(),
                'origin_type': request.form.get('tag_origin_type', '').strip(),
                'nsfw': 'tag_nsfw' in request.form,
            }

            form_data = {
                'name': name, 'character_id': character_id,
                'lora_lora_name': lora_name, 'lora_lora_weight': lora_weight,
                'lora_lora_triggers': lora_triggers, 'positive': positive,
                'negative': negative, 'tags': tags,
            }

            look_id = re.sub(r'[^a-zA-Z0-9_]', '_', name.lower().replace(' ', '_'))
            filename = f'{look_id}.json'
            file_path = os.path.join(app.config['LOOKS_DIR'], filename)

            data = {
                'look_id': look_id,
                'look_name': name,
                'character_id': character_id,
                'positive': positive,
                'negative': negative,
                'lora': {'lora_name': lora_name, 'lora_weight': lora_weight, 'lora_triggers': lora_triggers},
                'tags': tags
            }

            try:
                os.makedirs(app.config['LOOKS_DIR'], exist_ok=True)
                with open(file_path, 'w') as f:
                    json.dump(data, f, indent=2)

                slug = re.sub(r'[^a-zA-Z0-9_]', '', look_id)
                new_look = Look(look_id=look_id, slug=slug, filename=filename, name=name,
                                character_id=character_id, data=data,
                                is_nsfw=tags.get('nsfw', False))
                db.session.add(new_look)
                db.session.commit()

                flash(f'Look "{name}" created!')
                return redirect(url_for('look_detail', slug=slug))
            except Exception as e:
                logger.exception("Save error: %s", e)
                flash(f"Failed to create look: {e}")
                return render_template('looks/create.html', characters=characters, loras=loras, form_data=form_data)

        return render_template('looks/create.html', characters=characters, loras=loras, form_data=form_data)

    @app.route('/looks/bulk_create', methods=['POST'])
    def bulk_create_looks_from_loras():
        _s = Settings.query.first()
        lora_dir = ((_s.lora_dir_characters if _s else None) or app.config['LORA_DIR']).rstrip('/')
        _lora_subfolder = os.path.basename(lora_dir)
        if not os.path.exists(lora_dir):
            flash('Looks LoRA directory not found.', 'error')
            return redirect(url_for('looks_index'))

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

        system_prompt = load_prompt('look_system.txt')
        if not system_prompt:
            flash('Look system prompt file not found.', 'error')
            return redirect(url_for('looks_index'))

        looks_dir = app.config['LOOKS_DIR']

        for filename in os.listdir(lora_dir):
            if not filename.endswith('.safetensors'):
                continue

            name_base = filename.rsplit('.', 1)[0]
            look_id = re.sub(r'[^a-zA-Z0-9_]', '_', name_base.lower())
            look_name = re.sub(r'[^a-zA-Z0-9]+', ' ', name_base).title()

            json_filename = f"{look_id}.json"
            json_path = os.path.join(looks_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(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, look_id, look_name, json_path, looks_dir, _lora_subfolder, html_content, system_prompt, is_existing):
                def task_fn(job):
                    prompt = f"Create a look profile for a character appearance LoRA based on the filename: '{filename}'"
                    if html_content:
                        prompt += f"\n\nHere is descriptive text extracted from an associated HTML file:\n###\n{html_content[:3000]}\n###"

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

                    look_data['look_id'] = look_id
                    look_data['look_name'] = look_name

                    if 'lora' not in look_data:
                        look_data['lora'] = {}
                    look_data['lora']['lora_name'] = f"Illustrious/{_lora_subfolder}/{filename}"
                    if not look_data['lora'].get('lora_triggers'):
                        look_data['lora']['lora_triggers'] = name_base
                    if look_data['lora'].get('lora_weight') is None:
                        look_data['lora']['lora_weight'] = 0.8
                    if look_data['lora'].get('lora_weight_min') is None:
                        look_data['lora']['lora_weight_min'] = 0.7
                    if look_data['lora'].get('lora_weight_max') is None:
                        look_data['lora']['lora_weight_max'] = 1.0

                    os.makedirs(looks_dir, exist_ok=True)
                    with open(json_path, 'w') as f:
                        json.dump(look_data, f, indent=2)

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

            job = _enqueue_task(f"Create look: {look_name}", make_task(filename, name_base, look_id, look_name, json_path, looks_dir, _lora_subfolder, html_content, system_prompt, is_existing))
            job_ids.append(job['id'])

        if job_ids:
            def sync_task(job):
                sync_looks()
                job['result'] = {'synced': True}
            _enqueue_task("Sync looks 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)} look tasks ({skipped} skipped).')
        return redirect(url_for('looks_index'))

