app/services/retrieval.py

"""Retrieval grounding module for fetching similar game examples.

Handles four dataset schemas:
  1. Old dataset (games_dataset.json) — uses ``expected_output`` key
  2. New scavenger_hunt dataset — uses ``output`` key with ``tasks`` array
  3. New hide_and_seek dataset — uses ``output`` key with ``hiding_zones`` + ``play_area``
  4. New tag dataset — uses ``output`` key with ``arena`` + ``safe_zones`` + ``movement_features``
"""

import json
from pathlib import Path
from typing import Any


# ── Dataset source labels ──────────────────────────────────────────────────
DATASET_SOURCE_OLD = "old"
DATASET_SOURCE_SCOPE = "scavenger_hunt"  # new schema, scavenger hunt
DATASET_SOURCE_HIDE_SEEK = "hide_and_seek"  # new schema, hide & seek
DATASET_SOURCE_TAG = "tag"  # new schema, tag


def load_games_dataset(path: str) -> list[dict]:
    """Load the games dataset from JSON file.
    
    Args:
        path: Path to dataset JSON file
        
    Returns:
        List of game records
    """
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)


# ── Schema detection helpers ────────────────────────────────────────────────

def _detect_schema(record: dict) -> str:
    """Detect which schema a record uses based on its keys.
    
    Detection priority:
      1. ``expected_output`` key → old dataset
      2. ``output`` key with ``hiding_zones`` → hide_and_seek
      3. ``output`` key with ``arena`` → tag
      4. ``output`` key with ``tasks`` → scavenger_hunt (new schema)
      5. Fallback → old schema
    
    Returns one of ``DATASET_SOURCE_OLD``, ``DATASET_SOURCE_SCOPE``,
    ``DATASET_SOURCE_HIDE_SEEK``, or ``DATASET_SOURCE_TAG``.
    """
    # Old dataset uses expected_output, new datasets use output
    if "expected_output" in record:
        return DATASET_SOURCE_OLD

    output = record.get("output", {})
    if not output:
        return DATASET_SOURCE_OLD

    # Hide & seek has hiding_zones; tag has arena; scavenger hunt (new) has tasks
    if output.get("hiding_zones") is not None:
        return DATASET_SOURCE_HIDE_SEEK
    if output.get("arena") is not None:
        return DATASET_SOURCE_TAG
    if output.get("tasks") is not None:
        return DATASET_SOURCE_SCOPE
    return DATASET_SOURCE_OLD


# ── Old schema normalizer ───────────────────────────────────────────────────

def _normalize_old_record(record: dict) -> dict:
    """Normalize a record from the old ``games_dataset.json`` schema.
    
    Expected structure:
      {input: {game_type, location: {city, area, location_type},
               preferences: {duration_minutes, num_players, difficulty, age_group}},
       expected_output: {rules: [str], tasks: [{task_id, description, location_hint, points, ...}],
                         hints: [[str]], safety_flags: {is_safe, flags}},
       metadata: {quality_score, source, notes}}
    """
    inp = record.get("input", {})
    loc = inp.get("location", {})
    prefs = inp.get("preferences", {})
    out = record.get("expected_output", {})
    meta = record.get("metadata", {})

    tasks = out.get("tasks", [])
    rules = out.get("rules", [])
    hints = out.get("hints", [])
    sf = out.get("safety_flags", {})

    return {
        "id": record.get("id"),
        "dataset_source": DATASET_SOURCE_OLD,
        "game_type": inp.get("game_type", ""),
        "city": loc.get("city", ""),
        "area": loc.get("area", ""),
        "location_type": loc.get("location_type", ""),
        "duration_minutes": prefs.get("duration_minutes"),
        "num_players": prefs.get("num_players"),
        "difficulty": prefs.get("difficulty", ""),
        "age_group": prefs.get("age_group", ""),
        "landscape_tags": loc.get("landscape_tags", []),
        "theme": prefs.get("theme", ""),
        "mobility": prefs.get("mobility", "standard"),
        "allow_transport": prefs.get("allow_transport", False),
        "players_count": prefs.get("num_players"),
        "team_count": prefs.get("team_count", 1),
        "num_tasks": len(tasks),
        "task_ids": [t.get("task_id") for t in tasks],
        "num_rules": len(rules),
        "num_hints": len(hints),
        "is_safe": sf.get("is_safe", True),
        "safety_flags_list": sf.get("flags", []),
        "quality_score": meta.get("quality_score"),
        "source": meta.get("source"),
        "notes": meta.get("notes", ""),
        "rules": rules,
        "tasks": tasks,
        "hints": hints,
        # Hide & seek placeholders (not present in old schema)
        "hiding_zones": [],
        "play_area": {},
        "seeker_guidance": {},
        "hide_seek_scoring": {},
    }


# ── New schema normalizers ──────────────────────────────────────────────────

def _normalize_new_scavenger_record(record: dict) -> dict:
    """Normalize a record from the new scavenger_hunt dataset.
    
    Expected structure:
      {input: {game_type, location: {city, country, city_code, landscape_tags,
               urban_density, climate_zone, area_type},
               players: {count, team_count, age_group, mobility},
               preferences: {duration_minutes, difficulty, theme, allow_transport}},
       output: {game_title, rules: {objective, scoring_method, ...},
               tasks: [{task_id, title, description, landscape_tags_used,
                        task_type, difficulty_contribution, points,
                        completion_proof, estimated_time_minutes,
                        hints: {hint_1, hint_2, hint_3}, safety_flags}],
               safety_constraints: {...}, scoring_summary: {...}}}
    """
    inp = record.get("input", {})
    loc = inp.get("location", {})
    players = inp.get("players", {})
    prefs = inp.get("preferences", {})
    out = record.get("output", {})

    tasks = out.get("tasks", [])
    rules_obj = out.get("rules", {})
    safety = out.get("safety_constraints", {})

    # Flatten rules object into a list of rule strings
    rules_list = []
    if isinstance(rules_obj, dict):
        if rules_obj.get("objective"):
            rules_list.append(f"Objective: {rules_obj['objective']}")
        if rules_obj.get("team_rules"):
            rules_list.append(rules_obj["team_rules"])
        if rules_obj.get("disqualification_conditions"):
            for cond in rules_obj["disqualification_conditions"]:
                rules_list.append(f"No {cond}")
    elif isinstance(rules_obj, list):
        rules_list = rules_obj

    # Normalize each task: map new schema fields to unified keys
    normalized_tasks = []
    for t in tasks:
        nt = dict(t)
        # Map completion_proof → proof_type (new SH schema)
        if "completion_proof" in nt and "proof_type" not in nt:
            nt["proof_type"] = nt["completion_proof"]
        normalized_tasks.append(nt)

    return {
        "id": record.get("id"),
        "dataset_source": DATASET_SOURCE_SCOPE,
        "game_type": inp.get("game_type", "scavenger_hunt"),
        "city": loc.get("city", ""),
        "area": loc.get("area_type", loc.get("city", "")),
        "location_type": loc.get("area_type", ""),
        "duration_minutes": prefs.get("duration_minutes"),
        "num_players": players.get("count"),
        "difficulty": prefs.get("difficulty", ""),
        "age_group": players.get("age_group", ""),
        "landscape_tags": loc.get("landscape_tags", []),
        "theme": prefs.get("theme", ""),
        "mobility": players.get("mobility", "standard"),
        "allow_transport": prefs.get("allow_transport", False),
        "players_count": players.get("count"),
        "team_count": players.get("team_count", 1),
        "num_tasks": len(tasks),
        "task_ids": [t.get("task_id") for t in tasks],
        "num_rules": len(rules_list),
        "num_hints": sum(1 for t in tasks if t.get("hints")),
        "is_safe": True,
        "safety_flags_list": safety.get("exclusion_zones", []) + safety.get("physical_limits", []),
        "quality_score": out.get("quality_score"),
        "source": "new_scavenger",
        "notes": safety.get("notes", ""),
        "rules": rules_list,
        "tasks": normalized_tasks,  # normalized tasks with unified fields
        "hints": [t.get("hints", {}) for t in tasks],
        # Hide & seek placeholders
        "hiding_zones": [],
        "play_area": {},
        "seeker_guidance": {},
        "hide_seek_scoring": {},
    }


def _normalize_new_hide_seek_record(record: dict) -> dict:
    """Normalize a record from the new hide_and_seek dataset.
    
    Expected structure:
      {input: {game_type, location: {city, country, city_code, landscape_tags, ...},
               players: {count, team_count, age_group, mobility},
               preferences: {duration_minutes, difficulty, theme, allow_transport}},
       output: {game_title, rules: {objective, round_structure, round_count,
                hide_time_seconds, round_time_minutes, seeker_count, win_condition},
               play_area: {boundary_description, boundary_landscape_tags_used,
                           safe_base_description, boundary_size_tier, out_of_bounds},
               hiding_zones: [{zone_id, description, landscape_tags_used,
                               concealment_rating, risk_note, safety_flags}],
               seeker_guidance: {search_strategy, high_probability_zones},
               safety_constraints: {...}, scoring: {...}, ...}}
    """
    inp = record.get("input", {})
    loc = inp.get("location", {})
    players = inp.get("players", {})
    prefs = inp.get("preferences", {})
    out = record.get("output", {})

    rules_obj = out.get("rules", {})
    hiding_zones = out.get("hiding_zones", [])
    safety = out.get("safety_constraints", {})

    # Flatten rules
    rules_list = []
    if isinstance(rules_obj, dict):
        if rules_obj.get("objective"):
            rules_list.append(f"Objective: {rules_obj['objective']}")
        if rules_obj.get("round_structure"):
            rules_list.append(rules_obj["round_structure"])
        if rules_obj.get("win_condition"):
            rules_list.append(f"Win: {rules_obj['win_condition']}")
    elif isinstance(rules_obj, list):
        rules_list = rules_obj

    return {
        "id": record.get("id"),
        "dataset_source": DATASET_SOURCE_HIDE_SEEK,
        "game_type": inp.get("game_type", "hide_and_seek"),
        "city": loc.get("city", ""),
        "area": loc.get("area_type", loc.get("city", "")),
        "location_type": loc.get("area_type", ""),
        "duration_minutes": prefs.get("duration_minutes"),
        "num_players": players.get("count"),
        "difficulty": prefs.get("difficulty", ""),
        "age_group": players.get("age_group", ""),
        "landscape_tags": loc.get("landscape_tags", []),
        "theme": prefs.get("theme", ""),
        "mobility": players.get("mobility", "standard"),
        "allow_transport": prefs.get("allow_transport", False),
        "players_count": players.get("count"),
        "team_count": players.get("team_count", 1),
        "num_tasks": len(hiding_zones),
        "task_ids": [z.get("zone_id") for z in hiding_zones],
        "num_rules": len(rules_list),
        "num_hints": 0,
        "is_safe": True,
        "safety_flags_list": safety.get("exclusion_zones", []) + safety.get("physical_limits", []),
        "quality_score": out.get("quality_score"),
        "source": "new_hide_seek",
        "notes": safety.get("notes", ""),
        "rules": rules_list,
        "tasks": hiding_zones,  # hiding_zones stored as 'tasks' for unified retrieval
        "hints": [],
        # Hide & seek specific fields
        "hiding_zones": hiding_zones,
        "play_area": out.get("play_area", {}),
        "seeker_guidance": out.get("seeker_guidance", {}),
        "hide_seek_scoring": out.get("scoring", {}),
        # Tag placeholders
        "arena": {},
        "safe_zones": [],
        "movement_features": {},
    }


# ── New tag normalizer ────────────────────────────────────────────────────

def _normalize_new_tag_record(record: dict) -> dict:
    """Normalize a record from the new tag dataset.

    Expected structure:
      {input: {game_type, location: {city, ..., landscape_tags, ...},
               players: {count, team_count, age_group, mobility},
               preferences: {duration_minutes, difficulty, theme, allow_transport}},
       output: {game_title, rules: {objective, tag_variant, it_count, round_count, ...},
                arena: {boundary_description, boundary_size_tier, out_of_bounds},
                safe_zones: [{zone_id, description, capacity_rule, safety_flags}],
                movement_features: {chokepoints, open_zones, hazard_zones},
                safety_constraints: {...}, scoring: {...}, ...}}
    """
    inp = record.get("input", {})
    loc = inp.get("location", {})
    players = inp.get("players", {})
    prefs = inp.get("preferences", {})
    out = record.get("output", {})

    rules_obj = out.get("rules", {})
    safety = out.get("safety_constraints", {})
    arena = out.get("arena", {})
    safe_zones = out.get("safe_zones", [])

    # Flatten rules
    rules_list = []
    if isinstance(rules_obj, dict):
        if rules_obj.get("objective"):
            rules_list.append(f"Objective: {rules_obj['objective']}")
        if rules_obj.get("tag_mechanic"):
            rules_list.append(f"Mechanic: {rules_obj['tag_mechanic']}")
        if rules_obj.get("unfreeze_or_revival_rule"):
            rules_list.append(f"Revival: {rules_obj['unfreeze_or_revival_rule']}")
        if rules_obj.get("win_condition"):
            rules_list.append(f"Win: {rules_obj['win_condition']}")
    elif isinstance(rules_obj, list):
        rules_list = rules_obj

    return {
        "id": record.get("id"),
        "dataset_source": DATASET_SOURCE_TAG,
        "game_type": inp.get("game_type", "tag"),
        "city": loc.get("city", ""),
        "area": loc.get("area_type", loc.get("city", "")),
        "location_type": loc.get("area_type", ""),
        "duration_minutes": prefs.get("duration_minutes"),
        "num_players": players.get("count"),
        "difficulty": prefs.get("difficulty", ""),
        "age_group": players.get("age_group", ""),
        "landscape_tags": loc.get("landscape_tags", []),
        "theme": prefs.get("theme", ""),
        "mobility": players.get("mobility", "standard"),
        "allow_transport": prefs.get("allow_transport", False),
        "players_count": players.get("count"),
        "team_count": players.get("team_count", 1),
        "num_tasks": len(safe_zones),
        "task_ids": [z.get("zone_id") for z in safe_zones],
        "num_rules": len(rules_list),
        "num_hints": 0,
        "is_safe": True,
        "safety_flags_list": safety.get("exclusion_zones", []) + safety.get("physical_limits", []),
        "quality_score": out.get("quality_score"),
        "source": "new_tag",
        "notes": safety.get("notes", ""),
        "rules": rules_list,
        "tasks": safe_zones,
        "hints": [],
        # Tag specific fields
        "arena": arena,
        "safe_zones": safe_zones,
        "movement_features": out.get("movement_features", {}),
        "tag_variant": rules_obj.get("tag_variant", "classic"),
        "it_count": rules_obj.get("it_count", 1),
        "round_count": rules_obj.get("round_count", 1),
        "tag_mechanic": rules_obj.get("tag_mechanic", ""),
        "tag_scoring": out.get("scoring", {}),
        # Placeholders for other schemas
        "hiding_zones": [],
        "play_area": {},
        "seeker_guidance": {},
        "hide_seek_scoring": {},
    }


# ── Dispatcher ──────────────────────────────────────────────────────────────

def normalize_game_record(record: dict) -> dict:
    """Normalize a raw game record into structured format.
    
    Auto-detects the schema (old, new scavenger hunt, new hide & seek,
    or new tag) and dispatches to the appropriate normalizer.
    
    Args:
        record: Raw game record from any dataset
        
    Returns:
        Normalized game record with unified fields
    """
    schema = _detect_schema(record)
    if schema == DATASET_SOURCE_TAG:
        return _normalize_new_tag_record(record)
    if schema == DATASET_SOURCE_HIDE_SEEK:
        return _normalize_new_hide_seek_record(record)
    if schema == DATASET_SOURCE_SCOPE:
        return _normalize_new_scavenger_record(record)
    return _normalize_old_record(record)


# ── Similarity scoring ─────────────────────────────────────────────────────

def _compute_similarity_score(config: dict, record: dict) -> float:
    """Compute similarity score between a user config and a dataset record.
    
    Awards bonus points for:
      - Same game type (highest weight, +50)
      - New dataset source preferred (+15)
      - Age group match (+25/+15)
      - Difficulty match (+20)
      - Location type match (+15)
      - Landscape tag overlap (+10)
      - City name match (+10)
      - Duration proximity (+10/+5)
      - Theme match (+8)
      - Mobility match (+5)
      - Transport preference match (+5)
      - Quality bonus (+0.5 × score)
      - Safety bonus (+5)
    """
    score = 0.0

    # ── Game type exact match (highest weight) ───────────────────────
    if config.get('game_type', '').lower() == record.get('game_type', '').lower():
        score += 50

    # ── Dataset source bonus — prefer new, richer datasets ──────────
    src = record.get('dataset_source', DATASET_SOURCE_OLD)
    if src in (DATASET_SOURCE_SCOPE, DATASET_SOURCE_HIDE_SEEK, DATASET_SOURCE_TAG):
        score += 15

    # ── City match ──────────────────────────────────────────────────
    config_city = config.get('city', '').lower().strip()
    record_city = record.get('city', '').lower().strip()
    if config_city and record_city:
        if config_city == record_city:
            score += 10  # same city is a strong signal

    # ── Landscape tags overlap ──────────────────────────────────────
    config_tags = set(config.get('landscape_tags', []))
    record_tags = set(record.get('landscape_tags', []))
    if config_tags and record_tags:
        overlap = len(config_tags & record_tags)
        if overlap > 0:
            score += min(overlap * 3, 10)  # max +10

    # ── Age group match ─────────────────────────────────────────────
    config_age = config.get('age_group', '').lower()
    record_age = record.get('age_group', '').lower()
    if config_age and record_age:
        if config_age == record_age:
            score += 25
        elif 'mixed' in [config_age, record_age]:
            score += 15

    # ── Difficulty match ────────────────────────────────────────────
    if config.get('difficulty', '').lower() == record.get('difficulty', '').lower():
        score += 20

    # ── Location type match ─────────────────────────────────────────
    if config.get('location_type', '').lower() == record.get('location_type', '').lower():
        score += 15

    # ── Theme match ─────────────────────────────────────────────────
    config_theme = config.get('theme', '').lower().strip()
    record_theme = record.get('theme', '').lower().strip()
    if config_theme and record_theme and config_theme == record_theme:
        score += 8

    # ── Duration proximity ──────────────────────────────────────────
    config_duration = config.get('duration_minutes')
    record_duration = record.get('duration_minutes')
    if config_duration and record_duration:
        duration_diff = abs(config_duration - record_duration)
        if duration_diff <= 15:
            score += 10
        elif duration_diff <= 30:
            score += 5

    # ── Mobility match ──────────────────────────────────────────────
    config_mob = config.get('mobility', '').lower()
    record_mob = record.get('mobility', '').lower()
    if config_mob and record_mob and config_mob == record_mob:
        score += 5

    # ── Transport preference match ──────────────────────────────────
    config_transport = config.get('allow_transport')
    record_transport = record.get('allow_transport')
    if config_transport is not None and record_transport is not None:
        if config_transport == record_transport:
            score += 5

    # ── Quality bonus ───────────────────────────────────────────────
    quality = record.get('quality_score', 0)
    if quality:
        score += quality * 0.5

    # ── Safety bonus ────────────────────────────────────────────────
    if record.get('is_safe', True):
        score += 5

    return score


# ── Game-type-aware exemplar builder ───────────────────────────────────────

def _build_exemplar(record: dict, retrieval_score: float) -> dict:
    """Build a compressed exemplar bundle from a normalized record.
    
    The exemplar shape differs by game type so the generation prompt
    receives relevant context.
    """
    base = {
        'id': record['id'],
        'dataset_source': record.get('dataset_source', DATASET_SOURCE_OLD),
        'game_type': record['game_type'],
        'area': record.get('area', ''),
        'city': record.get('city', ''),
        'difficulty': record.get('difficulty', ''),
        'age_group': record.get('age_group', ''),
        'duration_minutes': record.get('duration_minutes'),
        'location_type': record.get('location_type', ''),
        'landscape_tags': record.get('landscape_tags', []),
        'theme': record.get('theme', ''),
        'mobility': record.get('mobility', 'standard'),
        'num_players': record.get('num_players') or record.get('players_count'),
        'rules_summary': record.get('rules', [])[:3],
        'safety_patterns': record.get('safety_flags_list', []),
        'notes': record.get('notes', '')[:120],
        'quality_score': record.get('quality_score'),
        'retrieval_score': retrieval_score,
    }

    game_type = record.get('game_type', '')

    if game_type == 'hide_and_seek':
        # Hide & seek: expose hiding zones, play area, seeker guidance
        base['hiding_zones_summary'] = [
            {
                'zone_id': z.get('zone_id'),
                'description': z.get('description', '')[:120],
                'concealment_rating': z.get('concealment_rating', 'medium'),
                'landscape_tags_used': z.get('landscape_tags_used', []),
            }
            for z in record.get('hiding_zones', [])[:3]
        ]
        base['play_area_summary'] = {
            k: record.get('play_area', {}).get(k, '')
            for k in ('boundary_description', 'boundary_size_tier', 'out_of_bounds')
            if record.get('play_area', {}).get(k)
        }
        base['seeker_strategy'] = (
            record.get('seeker_guidance', {}).get('search_strategy', '')[:200]
        )
        base['hide_seek_scoring'] = record.get('hide_seek_scoring', {})
        base['task_patterns'] = []

    elif game_type == 'tag' and record.get('dataset_source') == DATASET_SOURCE_TAG:
        # Tag (new schema): expose arena, safe zones, movement features
        arena = record.get('arena', {})
        base['arena_summary'] = {
            k: arena.get(k, '')
            for k in ('boundary_description', 'arena_size_tier', 'out_of_bounds')
            if arena.get(k)
        }
        base['safe_zones_summary'] = [
            {
                'zone_id': z.get('zone_id'),
                'description': z.get('description', '')[:120],
                'capacity_rule': z.get('capacity_rule', ''),
            }
            for z in record.get('safe_zones', [])[:3]
        ]
        mf = record.get('movement_features', {})
        base['chokepoints'] = mf.get('chokepoints', [])[:3]
        base['open_zones'] = mf.get('open_zones', [])[:3]
        base['tag_variant'] = record.get('tag_variant', 'classic')
        base['it_count'] = record.get('it_count', 1)
        base['round_count'] = record.get('round_count', 1)
        base['tag_mechanic'] = record.get('tag_mechanic', '')[:150]
        base['tag_scoring'] = record.get('tag_scoring', {})
        base['task_patterns'] = []

    else:
        # Scavenger hunt (old or new): expose task patterns
        tasks = record.get('tasks', [])
        base['task_patterns'] = [
            {
                'task_id': t.get('task_id'),
                'title': t.get('title', ''),
                'difficulty': t.get('difficulty_contribution', ''),
                'task_type': t.get('task_type', ''),
                'points': t.get('points'),
                'proof_type': t.get('proof_type', 'observation'),
                'time_limit': t.get('estimated_time_minutes') or t.get('time_limit_minutes'),
                'landscape_tags_used': t.get('landscape_tags_used', []),
            }
            for t in tasks[:3]
        ]

    # Ensure all exemplars have these keys (even if empty)
    base.setdefault('hiding_zones_summary', [])
    base.setdefault('task_patterns', [])

    return base


# ── Public retrieval API ────────────────────────────────────────────────────

def retrieve_examples(config: dict, dataset: list[dict], k: int = 3) -> list[dict]:
    """Retrieve top k closest examples from dataset.
    
    Uses a multi-factor similarity score (game type, city, tags, age group,
    difficulty, duration, theme, mobility) to rank all records, then returns
    game-type-aware compressed exemplar bundles.
    
    Args:
        config: Game configuration from user input
        dataset: Loaded and normalized dataset records (from all sources)
        k: Number of examples to retrieve
        
    Returns:
        List of retrieved example games (compressed exemplar bundles)
    """
    if not dataset:
        return []

    requested_type = config.get('game_type', '').lower()

    # Score all records
    scored_records = []
    for record in dataset:
        # Primary filter: must match game type (same-type examples are most useful)
        record_type = record.get('game_type', '').lower()
        if record_type != requested_type:
            continue

        score = _compute_similarity_score(config, record)
        scored_records.append({'record': record, 'score': score})

    # Sort by score descending
    scored_records.sort(key=lambda x: x['score'], reverse=True)

    # Take top k and build exemplar bundles
    top_k = scored_records[:k]
    return [_build_exemplar(item['record'], item['score']) for item in top_k]