src/combinations.py

"""Outfit combination engine: rules + LLM ranking.

Generates top+bottom outfit combinations from the catalog, filters by
compatibility rules (season, formality), and optionally ranks them
using the LLM. Persists user preferences (like/dislike) for future
reference.
"""

import json
import logging
import re
from datetime import datetime, timezone
from itertools import product
from pathlib import Path

from .catalog import load_catalog, get_garment_image_path
from .model_loader import model_manager

logger = logging.getLogger(__name__)

OUTFITS_PATH = Path(__file__).resolve().parent.parent / "data" / "outfits.json"

TOPS = frozenset({
    "shirt", "blouse", "t-shirt", "top", "tank-top",
    "sweater", "cardigan", "hoodie", "sweatshirt",
})

BOTTOMS = frozenset({
    "jeans", "pants", "trousers", "shorts", "skirt",
})

FORMALITY_COMPAT = {
    "casual": {"casual", "smart-casual"},
    "smart-casual": {"casual", "smart-casual", "formal"},
    "formal": {"smart-casual", "formal"},
}

RANKING_SYSTEM_PROMPT = (
    "You are a personal stylist. Rank outfit combinations (top + bottom) best-to-worst "
    "for the given occasion. Prioritise: occasion fit, color harmony, formality match, "
    "season, pattern balance. Return ONLY a JSON array of outfit IDs, best first."
)

# Max combos sent to the LLM — must fit in n_ctx=4096 alongside the response.
MAX_RANKING_ITEMS = 20


def _format_garment_line(garment: dict) -> str:
    """Compact one-line garment summary for the ranking prompt."""
    return (
        f"{garment.get('color', '?')} {garment.get('type', '?')}"
        f" ({garment.get('pattern', 'solid')}, {garment.get('season', 'all')},"
        f" {garment.get('formality', 'casual')})"
    )


def _select_combos_for_ranking(combinations: list[dict], max_items: int = MAX_RANKING_ITEMS) -> list[dict]:
    """Pick a diverse subset when the list is too large for the LLM context."""
    if len(combinations) <= max_items:
        return combinations

    # Spread selections across different tops so the LLM sees variety
    by_top: dict[str, list[dict]] = {}
    for combo in combinations:
        top_id = combo["top"]["id"]
        by_top.setdefault(top_id, []).append(combo)

    selected: list[dict] = []
    top_ids = list(by_top.keys())
    idx = 0
    while len(selected) < max_items and any(by_top.values()):
        top_id = top_ids[idx % len(top_ids)]
        bucket = by_top.get(top_id, [])
        if bucket:
            selected.append(bucket.pop(0))
        idx += 1

    return selected


def _format_liked_hint() -> str:
    """Build a short hint from previously liked outfits."""
    liked = get_liked_outfits()
    if not liked:
        return ""

    lines = ["\nUser liked (style signal):"]
    for outfit in liked[:3]:
        top = outfit["top"]
        bottom = outfit["bottom"]
        lines.append(
            f"- {_format_garment_line(top)} + {_format_garment_line(bottom)}"
        )
    return "\n".join(lines)


def _is_season_compatible(season_a: str, season_b: str) -> bool:
    if season_a == "all" or season_b == "all":
        return True
    return season_a == season_b


def _is_formality_compatible(formality_a: str, formality_b: str) -> bool:
    allowed = FORMALITY_COMPAT.get(formality_a, {formality_a})
    return formality_b in allowed


def _classify_garment(garment: dict) -> str | None:
    """Classify a garment as 'top', 'bottom', or None."""
    gtype = garment.get("type", "").lower().strip()

    if gtype in TOPS:
        return "top"
    for top_type in TOPS:
        if top_type in gtype or gtype in top_type:
            return "top"

    if gtype in BOTTOMS:
        return "bottom"
    for bottom_type in BOTTOMS:
        if bottom_type in gtype or gtype in bottom_type:
            return "bottom"

    return None


def get_tops_and_bottoms() -> tuple[list[dict], list[dict]]:
    """Split catalog into tops and bottoms."""
    catalog = load_catalog()
    tops = []
    bottoms = []

    for g in catalog:
        category = _classify_garment(g)
        if category == "top":
            tops.append(g)
        elif category == "bottom":
            bottoms.append(g)

    return tops, bottoms


def generate_combinations(
    season: str | None = None,
    exclude_disliked: bool = True,
) -> list[dict]:
    """Generate compatible top+bottom combinations.

    Applies rules-based filtering: season compatibility and formality
    compatibility. Optionally excludes previously disliked outfits.

    Returns a list of combination dicts:
    {"top": garment_dict, "bottom": garment_dict, "id": "outfit_NNN"}
    """
    tops, bottoms = get_tops_and_bottoms()

    if not tops or not bottoms:
        return []

    disliked_pairs = set()
    if exclude_disliked:
        disliked_pairs = _get_disliked_pairs()

    combinations = []
    combo_id = 1

    for top, bottom in product(tops, bottoms):
        pair_key = (top["id"], bottom["id"])
        if pair_key in disliked_pairs:
            continue

        top_season = top.get("season", "all")
        bottom_season = bottom.get("season", "all")
        if season:
            if not _is_season_compatible(top_season, season):
                continue
            if not _is_season_compatible(bottom_season, season):
                continue
        elif not _is_season_compatible(top_season, bottom_season):
            continue

        top_formality = top.get("formality", "casual")
        bottom_formality = bottom.get("formality", "casual")
        if not _is_formality_compatible(top_formality, bottom_formality):
            continue

        combinations.append({
            "id": f"outfit_{combo_id:03d}",
            "top": top,
            "bottom": bottom,
        })
        combo_id += 1

    return combinations


def rank_combinations_prompt(
    combinations: list[dict],
    context: str = "",
    max_items: int = MAX_RANKING_ITEMS,
) -> tuple[str, str]:
    """Build system + user prompts for the LLM to rank outfit combinations.

    Returns (system_prompt, user_prompt).
    """
    if not combinations:
        return "", ""

    subset = _select_combos_for_ranking(combinations, max_items=max_items)
    occasion = context.strip() if context and context.strip() else "everyday casual wear"

    user_lines = [
        f"Occasion: {occasion}",
        f"Rank these {len(subset)} outfits best-to-worst. Return JSON array of IDs only.",
        "",
    ]

    for combo in subset:
        top = _format_garment_line(combo["top"])
        bottom = _format_garment_line(combo["bottom"])
        user_lines.append(f"- {combo['id']}: {top} + {bottom}")

    liked_hint = _format_liked_hint()
    if liked_hint:
        user_lines.append(liked_hint)

    user_lines.append(f'Return: ["outfit_XXX", ...] with all {len(subset)} IDs reordered.')

    return RANKING_SYSTEM_PROMPT, "\n".join(user_lines)


def rank_with_llm(combinations: list[dict], context: str = "") -> list[dict]:
    """Rank combinations using the LLM.

    Always ranks — uses a default occasion when no context is provided.
    Falls back to the original order if parsing fails.
    """
    if not combinations:
        return combinations

    system_prompt, user_prompt = rank_combinations_prompt(combinations, context=context)
    if not user_prompt:
        return combinations

    llm = model_manager.get_text_model()

    logger.info(
        "Ranking %d combinations (context: %s)",
        len(combinations),
        (context[:80] if context else "default"),
    )

    response = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        max_tokens=512,
        temperature=0.2,
    )

    raw_text = response["choices"][0]["message"]["content"]
    logger.debug("LLM ranking response: %s", raw_text[:300])

    ranked_ids = _parse_ranking_response(raw_text)

    if not ranked_ids:
        logger.warning("Could not parse ranking, returning original order")
        return combinations

    combo_map = {c["id"]: c for c in combinations}
    ranked = [combo_map[oid] for oid in ranked_ids if oid in combo_map]

    seen = set(ranked_ids)
    for combo in combinations:
        if combo["id"] not in seen:
            ranked.append(combo)

    logger.info("Ranked %d/%d combinations successfully", len(ranked_ids), len(combinations))
    return ranked


def _parse_ranking_response(text: str) -> list[str]:
    """Extract a list of outfit IDs from the LLM ranking response."""
    cleaned = text.strip()

    fence_match = re.search(r"```(?:json)?\s*\n?(.*?)```", cleaned, re.DOTALL)
    if fence_match:
        cleaned = fence_match.group(1).strip()

    start = cleaned.find("[")
    end = cleaned.rfind("]")
    if start != -1 and end != -1 and end > start:
        try:
            parsed = json.loads(cleaned[start:end + 1])
            if isinstance(parsed, list):
                return [str(item) for item in parsed]
        except json.JSONDecodeError:
            pass

    ids = re.findall(r"outfit_\d+", cleaned)
    return ids


def load_outfits() -> list[dict]:
    """Load saved outfit preferences from disk."""
    if not OUTFITS_PATH.exists():
        return []
    try:
        with open(OUTFITS_PATH, "r", encoding="utf-8") as f:
            return json.load(f)
    except (json.JSONDecodeError, IOError) as e:
        logger.error("Failed to load outfits: %s", e)
        return []


def save_outfits(outfits: list[dict]) -> None:
    """Write outfit preferences to disk."""
    OUTFITS_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(OUTFITS_PATH, "w", encoding="utf-8") as f:
        json.dump(outfits, f, indent=2, ensure_ascii=False)


def save_preference(top_id: str, bottom_id: str, liked: bool) -> dict:
    """Record a user preference for an outfit combination."""
    outfits = load_outfits()

    for outfit in outfits:
        if outfit.get("top") == top_id and outfit.get("bottom") == bottom_id:
            outfit["liked"] = liked
            outfit["timestamp"] = datetime.now(timezone.utc).isoformat()
            save_outfits(outfits)
            return outfit

    next_num = len(outfits) + 1
    entry = {
        "id": f"outfit_{next_num:03d}",
        "top": top_id,
        "bottom": bottom_id,
        "liked": liked,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    outfits.append(entry)
    save_outfits(outfits)
    return entry


def get_liked_outfits() -> list[dict]:
    """Return only the outfits the user liked, with full garment data."""
    outfits = load_outfits()
    catalog = load_catalog()
    catalog_map = {g["id"]: g for g in catalog}

    liked = []
    for outfit in outfits:
        if not outfit.get("liked"):
            continue
        top = catalog_map.get(outfit["top"])
        bottom = catalog_map.get(outfit["bottom"])
        if top and bottom:
            liked.append({
                "id": outfit["id"],
                "top": top,
                "bottom": bottom,
                "timestamp": outfit.get("timestamp"),
            })

    return liked


def _get_disliked_pairs() -> set[tuple[str, str]]:
    """Return set of (top_id, bottom_id) pairs the user disliked."""
    outfits = load_outfits()
    return {
        (o["top"], o["bottom"])
        for o in outfits
        if not o.get("liked")
    }