# Copyright (C) 2025 AIDC-AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Character Memory component for web UI
"""

import streamlit as st
from typing import List

from web.i18n import tr
from web.utils.streamlit_helpers import safe_rerun
from pixelle_video.services.quality import CharacterMemory, Character, CharacterType


def render_active_characters_badge():
    """Render a small badge showing active characters count (for Home page)"""
    if "character_memory" not in st.session_state:
        return
    
    memory = st.session_state.character_memory
    active_count = sum(1 for c in memory.characters if getattr(c, 'is_active', True))
    
    if active_count > 0:
        st.info(f"👤 {tr('character.registered_count', count=active_count)} (Active)", icon="✨")


def render_character_memory():
    """Render the character memory management interface"""
    st.markdown(f"### {tr('character.title')}")
    st.caption(tr("character.desc"))
    
    # Initialize CharacterMemory in session state if not present
    if "character_memory" not in st.session_state:
        st.session_state.character_memory = CharacterMemory()
    
    memory = st.session_state.character_memory
    
    # --- Character Registration Form ---
    with st.expander(tr("character.add_button"), expanded=False):
        with st.form("add_character_form", clear_on_submit=True):
            col1, col2 = st.columns(2)
            
            with col1:
                name = st.text_input(tr("character.name"), key="new_char_name")
                char_type = st.selectbox(
                    tr("character.type"),
                    options=[t.value for t in CharacterType],
                    format_func=lambda x: tr(f"character.type.{x}"),
                    key="new_char_type"
                )
            
            with col2:
                features = st.text_input(tr("character.features"), placeholder="e.g., glasses, scar, red hat", key="new_char_features")
                # Image upload placeholder (Streamlit doesn't support persistent file paths easily in memory objects)
                ref_image = st.file_uploader(tr("character.ref_image"), type=["png", "jpg", "jpeg"], key="new_char_image")
            
            appearance = st.text_area(tr("character.appearance"), placeholder="e.g., young man with short black hair", key="new_char_appearance")
            clothing = st.text_area(tr("character.clothing"), placeholder="e.g., blue t-shirt and jeans", key="new_char_clothing")
            
            submit = st.form_submit_button(tr("character.add_button"), use_container_width=True)
            
            if submit:
                if not name:
                    st.error(tr("error.missing_field", field=tr("character.name")))
                else:
                    import os
                    import shutil
                    from pathlib import Path

                    feature_list = [f.strip() for f in features.split(",") if f.strip()]
                    
                    # Handle ref image upload
                    ref_path = None
                    if ref_image:
                        # Save to a temporary storage or artifacts dir
                        upload_dir = Path("/Users/yuanjiantsui/projects/Pixelle-Video/temp/characters")
                        upload_dir.mkdir(parents=True, exist_ok=True)
                        ref_path = str(upload_dir / f"{name}_{ref_image.name}")
                        with open(ref_path, "wb") as f:
                            f.write(ref_image.getbuffer())

                    char = memory.register_character(
                        name=name,
                        appearance_description=appearance,
                        clothing_description=clothing,
                        distinctive_features=feature_list,
                        character_type=CharacterType(char_type)
                    )
                    
                    if ref_path:
                        char.add_reference_image(ref_path, set_as_primary=True)
                        
                    st.success(f"✅ {name} registered!")
                    safe_rerun()

    # --- Character List ---
    st.markdown("---")
    characters = memory.characters
    
    if not characters:
        st.info(tr("character.no_characters"))
    else:
        st.markdown(tr("character.registered_count", count=len(characters)))
        
        # Display characters in a grid
        cols = st.columns(3)
        for i, char in enumerate(characters):
            with cols[i % 3]:
                with st.container(border=True):
                    # Header with Status Toggle
                    head_col, action_col = st.columns([3, 1])
                    with head_col:
                        st.markdown(f"**{char.name}**")
                        st.caption(tr(f"character.type.{char.character_type.value}"))
                    with action_col:
                        # Toggle for activation (defensive access for cached objects)
                        current_active = getattr(char, 'is_active', True)
                        new_state = st.toggle("On", value=current_active, key=f"tog_{char.id}", label_visibility="collapsed")
                        if new_state != current_active:
                            char.is_active = new_state
                            safe_rerun()
                    
                    # Image Preview
                    if char.primary_reference:
                        st.image(char.primary_reference, use_container_width=True)
                    else:
                        st.empty()
                    
                    if char.appearance_description:
                        st.caption(f"📝 {char.appearance_description}")
                    
                    if char.clothing_description:
                        st.caption(f"👕 {char.clothing_description}")
                        
                    if char.distinctive_features:
                        st.markdown(", ".join([f"`{f}`" for f in char.distinctive_features]))
                    
                    st.markdown("---")
                    # Delete Button at bottom
                    if st.button("🗑️ " + tr("history.task_card.delete"), key=f"del_{char.id}", use_container_width=True):
                        # Basic delete logic
                        del memory._characters[char.id]
                        if char.name.lower() in memory._name_index:
                            del memory._name_index[char.name.lower()]
                        safe_rerun()

    # Provide consistency prompt suggestion if characters exist
    if characters:
        with st.expander("💡 Character Consistency Prompt", expanded=False):
            st.code(
                ", ".join([char.get_prompt_injection() for char in characters]),
                language="text"
            )