character-browser/models.py

from flask_sqlalchemy import SQLAlchemy

db = SQLAlchemy()

class Character(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    character_id = db.Column(db.String(100), unique=True, nullable=False)
    slug = db.Column(db.String(100), unique=True, nullable=False)
    filename = db.Column(db.String(255), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    data = db.Column(db.JSON, nullable=False)
    default_fields = db.Column(db.JSON, nullable=True)
    image_path = db.Column(db.String(255), nullable=True)
    active_outfit = db.Column(db.String(100), default='default')
    
    # NEW: Outfit assignment support (Phase 4)
    assigned_outfit_ids = db.Column(db.JSON, default=list)  # List of outfit_ids from Outfit table
    default_outfit_id = db.Column(db.String(100), default='default')  # 'default' or specific outfit_id

    def get_active_wardrobe(self):
        """Get the currently active wardrobe outfit.
        
        Priority:
        1. If active_outfit is an outfit_id from Outfit table, fetch from Outfit.data['wardrobe']
        2. If active_outfit is a key in character's embedded wardrobe, use that
        3. Fall back to 'default'
        """
        # First check if active_outfit is an outfit_id from assigned outfits
        if self.active_outfit and self.active_outfit != 'default':
            # Try to get from Outfit table
            outfit = Outfit.query.filter_by(outfit_id=self.active_outfit).first()
            if outfit and outfit.data:
                return outfit.data.get('wardrobe', {})
        
        # Fall back to embedded wardrobe
        wardrobe = self.data.get('wardrobe', {})
        # Check if wardrobe is nested (new format) or flat (legacy)
        if 'default' in wardrobe and isinstance(wardrobe.get('default'), dict):
            # New nested format - return active outfit
            return wardrobe.get(self.active_outfit or 'default', wardrobe.get('default', {}))
        else:
            # Legacy flat format - return as-is
            return wardrobe
    
    def get_available_outfits(self):
        """Get list of available outfit objects (including embedded and assigned).
        
        Returns list of dicts with keys: outfit_id, name, source ('embedded' or 'assigned')
        """
        outfits = [{'outfit_id': 'default', 'name': 'Default', 'source': 'embedded'}]
        
        # Add embedded outfits from character data
        wardrobe = self.data.get('wardrobe', {})
        if 'default' in wardrobe and isinstance(wardrobe.get('default'), dict):
            for outfit_name in wardrobe.keys():
                if outfit_name != 'default':
                    outfits.append({
                        'outfit_id': outfit_name,
                        'name': outfit_name.replace('_', ' ').title(),
                        'source': 'embedded'
                    })
        
        # Add assigned outfits from Outfit table
        if self.assigned_outfit_ids:
            for outfit_id in self.assigned_outfit_ids:
                outfit = Outfit.query.filter_by(outfit_id=outfit_id).first()
                if outfit:
                    outfits.append({
                        'outfit_id': outfit.outfit_id,
                        'name': outfit.name,
                        'source': 'assigned'
                    })
        
        return outfits
    
    def get_outfit_wardrobe(self, outfit_id=None):
        """Get wardrobe data for a specific outfit.
        
        Args:
            outfit_id: Outfit ID to get wardrobe for. If None, uses active_outfit.
        
        Returns:
            Dict with wardrobe fields, or empty dict if not found.
        """
        if outfit_id is None:
            outfit_id = self.active_outfit or 'default'
        
        if outfit_id == 'default':
            # Return embedded default wardrobe
            wardrobe = self.data.get('wardrobe', {})
            if 'default' in wardrobe and isinstance(wardrobe.get('default'), dict):
                return wardrobe.get('default', {})
            return wardrobe
        
        # Try to find in Outfit table
        outfit = Outfit.query.filter_by(outfit_id=outfit_id).first()
        if outfit and outfit.data:
            return outfit.data.get('wardrobe', {})
        
        # Try embedded outfits
        wardrobe = self.data.get('wardrobe', {})
        if 'default' in wardrobe and isinstance(wardrobe.get('default'), dict):
            return wardrobe.get(outfit_id, {})
        
        return {}
    
    def assign_outfit(self, outfit_id):
        """Assign an outfit to this character.
        
        Args:
            outfit_id: The outfit_id from the Outfit table to assign.
        
        Returns:
            True if assigned, False if already assigned or outfit not found.
        """
        current_ids = self.assigned_outfit_ids or []
        
        # Verify outfit exists
        outfit = Outfit.query.filter_by(outfit_id=outfit_id).first()
        if not outfit:
            return False
        
        if outfit_id not in current_ids:
            new_ids = list(current_ids)
            new_ids.append(outfit_id)
            self.assigned_outfit_ids = new_ids
            return True
        return False
    
    def unassign_outfit(self, outfit_id):
        """Unassign an outfit from this character.
        
        Args:
            outfit_id: The outfit_id to unassign.
        
        Returns:
            True if unassigned, False if not found in assigned list.
        """
        current_ids = self.assigned_outfit_ids or []
        if outfit_id in current_ids:
            new_ids = list(current_ids)
            new_ids.remove(outfit_id)
            self.assigned_outfit_ids = new_ids
            # Reset active outfit if we just removed it
            if self.active_outfit == outfit_id:
                self.active_outfit = 'default'
            return True
        return False

    def __repr__(self):
        return f'<Character {self.character_id}>'

class Look(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    look_id = db.Column(db.String(100), unique=True, nullable=False)
    slug = db.Column(db.String(100), unique=True, nullable=False)
    filename = db.Column(db.String(255), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    character_id = db.Column(db.String(100), nullable=True)  # DEPRECATED: keeping for migration
    character_ids = db.Column(db.JSON, default=list)  # NEW: List of character_ids
    data = db.Column(db.JSON, nullable=False)
    default_fields = db.Column(db.JSON, nullable=True)
    image_path = db.Column(db.String(255), nullable=True)

    def get_linked_characters(self):
        """Get all characters linked to this look."""
        if not self.character_ids:
            return []
        return Character.query.filter(Character.character_id.in_(self.character_ids)).all()

    def add_character(self, character_id):
        """Link a character to this look."""
        if not self.character_ids:
            self.character_ids = []
        if character_id not in self.character_ids:
            self.character_ids.append(character_id)

    def remove_character(self, character_id):
        """Unlink a character from this look."""
        if self.character_ids and character_id in self.character_ids:
            self.character_ids.remove(character_id)

    def __repr__(self):
        return f'<Look {self.look_id}>'

class Outfit(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    outfit_id = db.Column(db.String(100), unique=True, nullable=False)
    slug = db.Column(db.String(100), unique=True, nullable=False)
    filename = db.Column(db.String(255), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    data = db.Column(db.JSON, nullable=False)
    default_fields = db.Column(db.JSON, nullable=True)
    image_path = db.Column(db.String(255), nullable=True)

    def __repr__(self):
        return f'<Outfit {self.outfit_id}>'

class Action(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    action_id = db.Column(db.String(100), unique=True, nullable=False)
    slug = db.Column(db.String(100), unique=True, nullable=False)
    filename = db.Column(db.String(255), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    data = db.Column(db.JSON, nullable=False)
    default_fields = db.Column(db.JSON, nullable=True)
    image_path = db.Column(db.String(255), nullable=True)

    def __repr__(self):
        return f'<Action {self.action_id}>'

class Style(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    style_id = db.Column(db.String(100), unique=True, nullable=False)
    slug = db.Column(db.String(100), unique=True, nullable=False)
    filename = db.Column(db.String(255), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    data = db.Column(db.JSON, nullable=False)
    default_fields = db.Column(db.JSON, nullable=True)
    image_path = db.Column(db.String(255), nullable=True)

    def __repr__(self):
        return f'<Style {self.style_id}>'

class Scene(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    scene_id = db.Column(db.String(100), unique=True, nullable=False)
    slug = db.Column(db.String(100), unique=True, nullable=False)
    filename = db.Column(db.String(255), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    data = db.Column(db.JSON, nullable=False)
    default_fields = db.Column(db.JSON, nullable=True)
    image_path = db.Column(db.String(255), nullable=True)

    def __repr__(self):
        return f'<Scene {self.scene_id}>'

class Detailer(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    detailer_id = db.Column(db.String(100), unique=True, nullable=False)
    slug = db.Column(db.String(100), unique=True, nullable=False)
    filename = db.Column(db.String(255), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    data = db.Column(db.JSON, nullable=False)
    default_fields = db.Column(db.JSON, nullable=True)
    image_path = db.Column(db.String(255), nullable=True)

    def __repr__(self):
        return f'<Detailer {self.detailer_id}>'

class Checkpoint(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    checkpoint_id = db.Column(db.String(255), unique=True, nullable=False)
    slug = db.Column(db.String(255), unique=True, nullable=False)
    name = db.Column(db.String(255), nullable=False)
    checkpoint_path = db.Column(db.String(255), nullable=False)  # e.g. "Illustrious/model.safetensors"
    data = db.Column(db.JSON, nullable=True)
    image_path = db.Column(db.String(255), nullable=True)

    def __repr__(self):
        return f'<Checkpoint {self.checkpoint_id}>'

class Preset(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    preset_id = db.Column(db.String(100), unique=True, nullable=False)
    slug = db.Column(db.String(100), unique=True, nullable=False)
    filename = db.Column(db.String(255), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    data = db.Column(db.JSON, nullable=False)
    image_path = db.Column(db.String(255), nullable=True)

    def __repr__(self):
        return f'<Preset {self.preset_id}>'


class Settings(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    llm_provider = db.Column(db.String(50), default='openrouter') # 'openrouter', 'ollama', 'lmstudio'
    openrouter_api_key = db.Column(db.String(255), nullable=True)
    openrouter_model = db.Column(db.String(100), default='google/gemini-2.0-flash-001')
    local_base_url = db.Column(db.String(255), nullable=True)
    local_model = db.Column(db.String(100), nullable=True)
    # LoRA directories (absolute paths on disk)
    lora_dir_characters = db.Column(db.String(500), default='/ImageModels/lora/Illustrious/Looks')
    lora_dir_outfits    = db.Column(db.String(500), default='/ImageModels/lora/Illustrious/Clothing')
    lora_dir_actions    = db.Column(db.String(500), default='/ImageModels/lora/Illustrious/Poses')
    lora_dir_styles     = db.Column(db.String(500), default='/ImageModels/lora/Illustrious/Styles')
    lora_dir_scenes     = db.Column(db.String(500), default='/ImageModels/lora/Illustrious/Backgrounds')
    lora_dir_detailers  = db.Column(db.String(500), default='/ImageModels/lora/Illustrious/Detailers')
    # Checkpoint scan directories (comma-separated list of absolute paths)
    checkpoint_dirs     = db.Column(db.String(1000), default='/ImageModels/Stable-diffusion/Illustrious,/ImageModels/Stable-diffusion/Noob')
    # Default checkpoint path (persisted across server restarts)
    default_checkpoint  = db.Column(db.String(500), nullable=True)
    # API key for REST API authentication
    api_key = db.Column(db.String(255), nullable=True)

    def __repr__(self):
        return '<Settings>'