import os
import torch
import pandas as pd
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load model
model_name = "google/flan-t5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Load all product CSVs from nested directories
def load_all_product_data(base_path="data"):
    all_data = []
    for root, dirs, files in os.walk(base_path):
        for file in files:
            if file.endswith(".csv"):
                full_path = os.path.join(root, file)
                df = pd.read_csv(full_path)
                df["Brand"] = os.path.splitext(file)[0]
                df["Store"] = root.split(os.sep)[-4]  # e.g., "Store C"
                df["Category"] = root.split(os.sep)[-2]
                all_data.append(df)
    return pd.concat(all_data, ignore_index=True)

df = load_all_product_data("data")

# Generate smart suggestion
def generate_product_description(prompt):
    prompt = prompt.lower()

    # Basic price filter
    price_limit = 99999
    if "under" in prompt:
        try:
            price_limit = int(prompt.split("under")[-1].split()[0])
        except:
            pass

    filtered_df = df[df["Price"] <= price_limit]
    filtered_df = filtered_df[df["In Stock"].str.lower() == "yes"]

    if "dry hair" in prompt:
        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("dry", na=False)]
    elif "oily hair" in prompt:
        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("oily", na=False)]
    elif "normal hair" in prompt:
        filtered_df = filtered_df[filtered_df["Hair Type"].str.lower().str.contains("normal", na=False)]

    if "gift" in prompt:
        filtered_df = filtered_df[filtered_df["Tags"].str.contains("gift", case=False, na=False)]
    if "budget" in prompt:
        filtered_df = filtered_df[filtered_df["Tags"].str.contains("budget", case=False, na=False)]

    if filtered_df.empty:
        return "🤷 Sorry, no matching suggestions found."

    rows = []
    for _, row in filtered_df.iterrows():
        line = f"{row['Brand']} {row['Quantity']} – ₹{row['Price']} (Floor {row['Floor']}, Aisle {row['Aisle']})"
        if pd.notna(row.get("Offer")) and str(row["Offer"]).strip():
            line += f"\n🎉 {row['Offer']}"
        rows.append(line)

    product_text = "\n".join(rows)
    model_prompt = f"Suggest top products:\n{product_text}"

    input_ids = tokenizer(model_prompt, return_tensors="pt").input_ids
    with torch.no_grad():
        output_ids = model.generate(input_ids, max_new_tokens=100)

    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)

    # Combine product suggestions and model response line by line
    return "\n\n".join(rows[:5]) + "\n\n🧠 AI Suggestion: " + response