import gradio as gr
from PIL import Image
import torch
import io
from transformers import (
    BlipProcessor,
    BlipForConditionalGeneration,
    CLIPProcessor,
    CLIPModel,
    AutoTokenizer,
    AutoModelForCausalLM,
)
import os
import json

# Global state for models
blip_model = None
blip_processor = None
clip_model = None
clip_processor = None
llm_tokenizer = None
llm_model = None

def device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'

def load_models():
    global blip_model, blip_processor, clip_model, clip_processor, llm_model, llm_tokenizer

    if blip_model is None:
        print('Loading BLIP image captioning model...')
        blip_processor = BlipProcessor.from_pretrained('Salesforce/blip-image-captioning-base')
        blip_model = BlipForConditionalGeneration.from_pretrained('Salesforce/blip-image-captioning-base').to(device())
        blip_model.eval()

    if clip_model is None:
        print('Loading CLIP model...')
        clip_processor = CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32')
        clip_model = CLIPModel.from_pretrained('openai/clip-vit-base-patch32').to(device())
        clip_model.eval()

    if os.environ.get('USE_TINY_LLAMA', '1') == '1' and llm_model is None:
        try:
            print('Loading TinyLlama chat model (CPU)...')
            llm_tokenizer = AutoTokenizer.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v1.0')
            llm_model = AutoModelForCausalLM.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v1.0').to('cpu')
            llm_model.eval()
        except Exception as e:
            print(f'Failed to load TinyLlama: {e}')

def analyze_fashion_image(pil_image: Image.Image):
    load_models()
    # BLIP caption
    inputs = blip_processor(pil_image, return_tensors='pt').to(device())
    with torch.no_grad():
        out = blip_model.generate(**inputs, max_length=64)
    caption = blip_processor.decode(out[0], skip_special_tokens=True)

    # Lightweight fashion prompts via BLIP conditional generation
    fashion_prompts = [
        'Describe the clothing type and key attributes.',
        'List the dominant colors.',
        'What style or vibe does this outfit convey?',
        'Suggest an occasion where this outfit fits well.'
    ]
    analysis = {}
    for prompt in fashion_prompts:
        try:
            cond_inputs = blip_processor(pil_image, prompt, return_tensors='pt').to(device())
            with torch.no_grad():
                out2 = blip_model.generate(**cond_inputs, max_length=48)
            resp = blip_processor.decode(out2[0], skip_special_tokens=True)
            analysis[prompt] = resp
        except Exception:
            analysis[prompt] = ''

    return {
        'caption': caption,
        'fashion_analysis': analysis,
    }

def simple_rule_based_advice(text: str) -> str:
    t = (text or '').lower()
    advice = []
    if 'black' in t:
        advice.append('Black is versatile; add a contrasting accessory (silver, gold, or a bold color) to elevate the look.')
    if 'white' in t:
        advice.append('White gives a clean aesthetic; consider layering with textured fabrics or a light jacket for depth.')
    if 'blue' in t:
        advice.append('Blue pairs well with neutrals (tan, white, grey). Try denim-on-denim or navy with beige for a classic combo.')
    if 'dress' in t:
        advice.append('Consider heel height and bag size to match the dress formality. A belt can define the silhouette if needed.')
    if 'shirt' in t or 'top' in t:
        advice.append('Balance the silhouette: pair fitted tops with relaxed bottoms or vice versa. Tuck/half-tuck for shape.')
    if not advice:
        advice.append('Consider color harmony, silhouette balance, and occasion. Add one statement accessory to complete the look.')
    return ' '.join(advice)

def llm_advise(system_prompt: str, user_prompt: str) -> str:
    # Try TinyLlama, else fallback to rule-based
    if llm_model is None or llm_tokenizer is None:
        return simple_rule_based_advice(user_prompt)
    prompt = (
        f"<|system|>\n{system_prompt}\n<|user|>\n{user_prompt}\n<|assistant|>"
    )
    inputs = llm_tokenizer(prompt, return_tensors='pt')
    with torch.no_grad():
        out = llm_model.generate(
            **inputs,
            max_new_tokens=220,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            eos_token_id=llm_tokenizer.eos_token_id,
        )
    text = llm_tokenizer.decode(out[0], skip_special_tokens=True)
    # Extract assistant part
    if '<|assistant|>' in text:
        text = text.split('<|assistant|>')[-1].strip()
    return text

def analyze_fashion(image):
    if image is None:
        return "Please upload an image first!"
    analysis = analyze_fashion_image(image)
    sys_prompt = (
        'You are a professional fashion stylist. Provide constructive, specific, and friendly advice about fit, colors, style, and accessories.'
    )
    user_msg = f"Analyze and advise based on: {json.dumps(analysis)}"
    advice = llm_advise(sys_prompt, user_msg)
    return advice

# --- Try-On API Integration ---
import requests
def tryon_api(clothing_img, avatar_img):
    clothing_bytes = io.BytesIO()
    clothing_img.save(clothing_bytes, format='JPEG')
    clothing_bytes.seek(0)
    avatar_bytes = io.BytesIO()
    avatar_img.save(avatar_bytes, format='JPEG')
    avatar_bytes.seek(0)
    files = {
        'clothing_image': ('clothing.jpg', clothing_bytes, 'image/jpeg'),
        'avatar_image': ('avatar.jpg', avatar_bytes, 'image/jpeg'),
    }
    headers = {
        'X-RapidAPI-Key': '76cebf2c24msh8ea6e489fb837f0p171e40jsn2f57db29afd2',
        'X-RapidAPI-Host': 'try-on-diffusion.p.rapidapi.com'
    }
    response = requests.post('https://try-on-diffusion.p.rapidapi.com/try-on-file', files=files, headers=headers)
    if response.status_code == 200:
        return Image.open(io.BytesIO(response.content))
    else:
        raise Exception('Try-on generation failed: ' + response.text)
    
import asyncio
import re
try:
    from playwright.async_api import async_playwright
except ImportError:
    async_playwright = None

FASHION_SITES = {
    'zara': {
        'url': 'https://www.zara.com/us/en/search',
        'search_param': 'searchTerm',
        'selectors': {
            'products': '.product-item',
            'title': '.product-link',
            'price': '.price',
            'image': '.media-image img',
            'link': '.product-link'
        }
    },
    'hm': {
        'url': 'https://www2.hm.com/en_us/search-results.html',
        'search_param': 'q',
        'selectors': {
            'products': '.item-link',
            'title': '.item-heading',
            'price': '.item-price',
            'image': '.item-image img',
            'link': '.item-link'
        }
    },
    'asos': {
        'url': 'https://www.asos.com/us/search/',
        'search_param': 'q',
        'selectors': {
            'products': '[data-testid="product-tile"]',
            'title': '[data-testid="product-title"]',
            'price': '[data-testid="current-price"]',
            'image': 'img',
            'link': 'a'
        }
    }
}

async def search_fashion_items(query, max_results=3):
    results = []
    if async_playwright is None:
        return results
    try:
        async with async_playwright() as p:
            browser = await p.chromium.launch(headless=True)
            context = await browser.new_context(
                user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
            )
            for site_name, site_config in FASHION_SITES.items():
                try:
                    page = await context.new_page()
                    search_url = f"{site_config['url']}?{site_config['search_param']}={query}"
                    await page.goto(search_url, wait_until='networkidle', timeout=10000)
                    await page.wait_for_timeout(2000)
                    products = await page.query_selector_all(site_config['selectors']['products'])
                    site_results = []
                    for product in products[:max_results]:
                        try:
                            title_elem = await product.query_selector(site_config['selectors']['title'])
                            price_elem = await product.query_selector(site_config['selectors']['price'])
                            image_elem = await product.query_selector(site_config['selectors']['image'])
                            link_elem = await product.query_selector(site_config['selectors']['link'])
                            title = await title_elem.inner_text() if title_elem else 'N/A'
                            price = await price_elem.inner_text() if price_elem else 'N/A'
                            image_src = await image_elem.get_attribute('src') if image_elem else ''
                            link_href = await link_elem.get_attribute('href') if link_elem else ''
                            title = title.strip()[:100]
                            price = re.sub(r'[^\\d.,$€£]', '', price) if price != 'N/A' else 'N/A'
                            if link_href and not link_href.startswith('http'):
                                base_url = f"https://{site_name}.com" if site_name != 'hm' else 'https://www2.hm.com'
                                link_href = base_url + link_href
                            site_results.append({
                                'title': title,
                                'price': price,
                                'image': image_src,
                                'link': link_href,
                                'store': site_name.upper(),
                                'query': query
                            })
                        except Exception:
                            pass
                    results.extend(site_results)
                    await page.close()
                except Exception:
                    pass
            await browser.close()
    except Exception:
        pass
    return results

def shopping_search_sync(query):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        results = loop.run_until_complete(search_fashion_items(query, max_results=3))
    finally:
        loop.close()
    if not results:
        return "No products found. Try a different keyword."
    out = ""
    for r in results:
        out += f"\n**{r['title']}**\nStore: {r['store']}\nPrice: {r['price']}\n[View Product]({r['link']})\n\n"
    return out

# --- Gradio UI: Try-On Only ---
with gr.Blocks() as demo:
    gr.Markdown("# 👗 AI Virtual Try-On\nUpload clothing and avatar images, generate a try-on result!")
    with gr.Row():
        clothing_input = gr.Image(type="pil", label="Clothing Image")
        avatar_input = gr.Image(type="pil", label="Avatar Image")

    tryon_btn = gr.Button("Generate Try-On & Analyze")
    tryon_output = gr.Image(label="Try-On Result")
    # analysis_output = gr.Textbox(label="AI Analysis")
    advice_output = gr.Textbox(label="AI Advice")


    # Store the latest analysis for chatbot context
    latest_analysis = {"caption": "", "fashion_analysis": {}}

    def tryon_workflow(cloth, avatar):
        img = tryon_api(cloth, avatar)
        analysis = analyze_fashion_image(img)
        # Save for chatbot context
        latest_analysis["caption"] = analysis["caption"]
        latest_analysis["fashion_analysis"] = analysis["fashion_analysis"]
        sys_prompt = 'You are a professional fashion stylist. Provide constructive, specific, and friendly advice about fit, colors, style, and accessories.'
        user_msg = f"Analyze and advise based on: {json.dumps(analysis)}"
        advice = llm_advise(sys_prompt, user_msg)
        pretty = f"📝 Caption: {analysis['caption']}\n"
        for q, a in analysis['fashion_analysis'].items():
            pretty += f"\n{q}\n➡️ {a}\n"
        return img, pretty, advice

    tryon_btn.click(tryon_workflow, inputs=[clothing_input, avatar_input], outputs=[tryon_output, advice_output])

    # --- Recommend Similar Outfits ---
    def recommend_similar(pil_image):
        load_models()
        inputs = clip_processor(images=pil_image, return_tensors='pt').to(device())
        with torch.no_grad():
            _ = clip_model.get_image_features(**inputs)
        recs = [
            {
                'title': 'Similar silhouette',
                'description': 'Find items with a matching cut and color palette.',
                'stores': ['Zara', 'H&M', 'ASOS']
            },
            {
                'title': 'Complementary accessories',
                'description': 'Belts, bags, and shoes that elevate this outfit.',
                'stores': ['Uniqlo', 'Mango', 'Amazon']
            },
            {
                'title': 'Alternative colors',
                'description': 'Same style in seasonal colorways to suit your palette.',
                'stores': ['Nordstrom', "Macy's", 'Urban Outfitters']
            }
        ]
        rec_text = 'Here are some similar outfit ideas and recommendations:\n'
        for idx, rec in enumerate(recs):
            rec_text += f"\n{idx + 1}. {rec['title']}: {rec['description']} (Stores: {', '.join(rec['stores'])})"
        return rec_text

    # --- Chatbot ---
    def chatbot(message, history=None):
        sys_prompt = 'You are a helpful, up-to-date fashion stylist assistant.'
        # Add latest analysis as context if available
        context = ""
        if latest_analysis["caption"] or latest_analysis["fashion_analysis"]:
            context = f"\n\nCurrent try-on analysis: {json.dumps(latest_analysis)}"
        return llm_advise(sys_prompt, message + context)

    # --- Gradio UI: Tabs for all features ---
    with gr.Tab("Chatbot"):
        chatbox = gr.ChatInterface(chatbot, title="Fashion Chatbot (TinyLlama)")

    with gr.Tab("Find Similar"):
        sim_input = gr.Image(type="pil", label="Upload Try-On or Fashion Image")
        sim_btn = gr.Button("Find Similar Outfits")
        sim_output = gr.Textbox(label="Similar Outfits & Recommendations")
        sim_btn.click(recommend_similar, inputs=sim_input, outputs=sim_output)
    
    with gr.Tab("Shopping Automation"):
        shop_query = gr.Textbox(label="Describe the item or keywords to search (e.g. 'black dress', 'summer shirt')")
        shop_btn = gr.Button("Search Fashion Stores")
        shop_results = gr.Markdown(label="Shopping Results")
        shop_btn.click(shopping_search_sync, inputs=shop_query, outputs=shop_results)

demo.launch()