wardrobe.py

import re
from typing import List, Optional
from models import WardrobeItem, ChatResponse
from query_processing import get_color_matches
from model_manager import generate_chat_response
from conversation import get_conversation_context, enhance_message_with_context, update_context
from rag import retrieve_relevant_context, format_rag_context

def clean_wardrobe_response(text: str) -> str:
    text = re.sub(r'\b(\w+)\s+\1\b', r'\1', text)
    text = re.sub(r'\b(navy|grey|gray|black|white|brown|blue|red|green|beige|tan|charcoal|rolex|fossil|hermes|zara|nike)\s+\1\b', r'\1', text, flags=re.IGNORECASE)
    
    if "protein" in text.lower():
        text = re.sub(r'[Pp]rotein\s*', '', text)
    
    lines = text.split('\n')
    cleaned_lines = []
    
    for line in lines:
        line = line.strip()
        if not line or len(line) < 3:
            continue
        if "protein" in line.lower():
            continue
        cleaned_lines.append(line)
    
    result = '\n'.join(cleaned_lines).strip()
    
    if len(result) > 1000:
        sentences = result.split('. ')
        result = '. '.join(sentences[:8]) + '.'
    
    return result

def format_wardrobe_for_prompt(wardrobe: List[WardrobeItem]) -> str:
    wardrobe_by_category = {}
    for item in wardrobe:
        category = item.category.lower()
        if category not in wardrobe_by_category:
            wardrobe_by_category[category] = []
        wardrobe_by_category[category].append(item)
    
    wardrobe_details = []
    for idx, item in enumerate(wardrobe, 1):
        parts = []
        if item.brand:
            parts.append(item.brand)
        if item.color:
            parts.append(item.color)
        if item.name:
            parts.append(item.name)
        elif item.category:
            parts.append(item.category)
        
        item_name = " ".join(parts) if parts else item.category
        wardrobe_details.append(f'{idx}. {item_name} ({item.category}, {item.style})')
    
    categories_list = ", ".join(wardrobe_by_category.keys())
    
    return f"""Available items ({len(wardrobe)} total):
{chr(10).join(wardrobe_details)}

Categories: {categories_list}"""

async def handle_wardrobe_chat(message: str, wardrobe: List[WardrobeItem], session_id: str, images: Optional[List[str]] = None, wardrobe_description: Optional[str] = None) -> ChatResponse:
    conv_context = get_conversation_context(session_id)
    enhanced_message = enhance_message_with_context(message, conv_context["context"])
    
    wardrobe_context = wardrobe_description if wardrobe_description else format_wardrobe_for_prompt(wardrobe)
    
    wardrobe_by_category = {}
    for item in wardrobe:
        category = item.category.lower()
        if category not in wardrobe_by_category:
            wardrobe_by_category[category] = []
        wardrobe_by_category[category].append(item)
    
    rag_chunks = retrieve_relevant_context(enhanced_message, top_k=3)
    rag_context = format_rag_context(rag_chunks)
    
    occasion_keywords = ["defense", "project", "presentation", "meeting", "interview", "formal", "casual", "party", "wedding", "dinner", "date", "work", "office"]
    occasion = next((word for word in occasion_keywords if word in enhanced_message.lower()), None)
    
    context_info = f"Available wardrobe categories: {', '.join(wardrobe_by_category.keys())}. "
    if occasion:
        context_info += f"Occasion: {occasion}. "
    
    system_override = "You are a friendly and helpful fashion stylist. Suggest complete outfits conversationally and warmly. Include accessories when available. Be natural and friendly in your responses."
    
    prompt = f"""{wardrobe_context}

User request: {enhanced_message}

Suggest a complete outfit using ONLY the items listed above. Reference items by their exact names as shown (e.g., if item is "zara black pants", say "zara black pants", not "black zara pants"). Include accessories (watches, bags, jewelry, belts, glasses) if available. Be friendly and conversational. Suggest: one top/shirt, one bottom (pants/shorts), shoes, and accessories. Explain briefly why it works."""
    
    if context_info.strip():
        prompt += f"\n\nContext: {context_info.strip()}"
    
    response_text = generate_chat_response(prompt, max_length=512, temperature=0.8, rag_context=rag_context, system_override=system_override, images=images)
    
    response_text = clean_wardrobe_response(response_text)
    
    update_context(session_id, message, {
        "response": response_text,
        "wardrobe_count": len(wardrobe),
        "categories": list(wardrobe_by_category.keys())
    })
    
    return ChatResponse(
        response=response_text,
        session_id=session_id
    )