# handlers.py
import os
import json
import asyncio
import re

import numpy as np
import pandas as pd
import psutil
import tskit

from lorax.modes import CURRENT_MODE
from lorax.cloud.gcs_utils import get_public_gcs_dict
from lorax.tree_graph import construct_trees_batch, construct_tree, TreeGraph
from lorax.csv.layout import build_empty_layout_response, build_csv_layout_response
from lorax.utils import (
    ensure_json_dict,
    list_project_files,
    make_json_serializable,
)
from lorax.metadata.loader import (
    get_metadata_for_key,
    search_samples_by_metadata,
    get_metadata_array_for_key,
    _get_sample_metadata_value
)
from lorax.metadata.mutations import (
    get_mutations_in_window,
    search_mutations_by_position
)
from lorax.buffer import mutations_to_arrow_buffer
from lorax.cache import get_file_context, get_file_cache_size
def _get_tip_shift_project_prefixes() -> list[str]:
    """Return project name prefixes that should shift CSV tips to y=1."""
    raw = os.getenv("LORAX_CSV_TIP_SHIFT_PROJECTS", "heliconius")
    parts = [p.strip().lower() for p in raw.split(",") if p.strip()]
    return parts


def should_shift_csv_tips(file_path: str) -> bool:
    """Return True when the file path matches a configured project prefix."""
    if not file_path:
        return False
    prefixes = _get_tip_shift_project_prefixes()
    if not prefixes:
        return False
    parts = re.split(r"[\\/]", str(file_path))
    for part in parts:
        if not part:
            continue
        part_l = str(part).lower()
        if any(part_l.startswith(prefix) for prefix in prefixes):
            return True
    return False


def _is_heliconius_project(file_path: str) -> bool:
    """Return True when the file path indicates a Heliconius project (case-insensitive)."""
    if not file_path:
        return False
    parts = re.split(r"[\\/]", str(file_path))
    for part in parts:
        if part and str(part).lower().startswith("heliconius"):
            return True
    return False


async def cache_status():
    """Return current memory usage and cache statistics."""
    process = psutil.Process(os.getpid())
    mem_info = process.memory_info()
    rss_mb = mem_info.rss / (1024 * 1024)
    vms_mb = mem_info.vms / (1024 * 1024)

    return {
        "rss_MB": round(rss_mb, 2),
        "vms_MB": round(vms_mb, 2),
        "file_cache_size": get_file_cache_size(),
        "pid": os.getpid(),
    }


async def handle_upload(file_path, root_dir):
    """Load a file and return its FileContext."""
    ctx = await get_file_context(file_path, root_dir)
    print("File loading complete")
    return ctx


async def get_projects(upload_dir, BUCKET_NAME, sid=None):
    """List all projects and their files from local uploads and GCS bucket."""
    projects = {}
    upload_dir = str(upload_dir)
    # Avoid listing Uploads/<sid> as separate projects; add session-scoped uploads below.
    projects = list_project_files(
        upload_dir,
        projects,
        root=upload_dir,
        exclude_dirs=["Uploads"],
    )

    # Prefer session-scoped Uploads/<sid> when available (non-local); local uses flat Uploads
    upload_files = []
    uploads_root = None
    if CURRENT_MODE == "local":
        uploads_root = os.path.join(upload_dir, "Uploads")
    else:
        uploads_root = os.path.join(upload_dir, "Uploads", sid) if sid else None

    if uploads_root and os.path.isdir(uploads_root):
        for item in os.listdir(uploads_root):
            if item.endswith((".trees", ".trees.tsz", ".csv")):
                upload_files.append(item)

    projects["Uploads"] = {
        "folder": "Uploads",
        "files": sorted(set(upload_files)),
        "description": "",
    }
    # Remove accidental project entry created from Uploads/<sid>
    if sid and sid in projects:
        projects.pop(sid, None)

    # Merge GCS projects: always include non-Uploads; Uploads only per mode rules
    if CURRENT_MODE == "local":
        projects = get_public_gcs_dict(
            BUCKET_NAME,
            sid=sid,
            projects=projects,
            include_uploads=False,
            uploads_sid=None,
        )
    else:
        projects = get_public_gcs_dict(
            BUCKET_NAME,
            sid=sid,
            projects=projects,
            include_uploads=True,
            uploads_sid=sid,
        )

    return projects

def _build_sample_name_mapping(ts, sample_name_key="name"):
    """
    Build mapping from sample name (lowercase) to node_id.

    Args:
        ts: tskit.TreeSequence
        sample_name_key: Key in node metadata used as sample name

    Returns:
        dict mapping lowercase sample name to node_id
    """
    name_to_node_id = {}
    for node_id in ts.samples():
        node = ts.node(node_id)
        node_meta = node.metadata or {}
        try:
            node_meta = ensure_json_dict(node_meta)
        except (TypeError, json.JSONDecodeError):
            node_meta = {}
        name = str(node_meta.get(sample_name_key, f"{node_id}"))
        name_to_node_id[name.lower()] = node_id
    return name_to_node_id


def _compute_lineage_paths(tree, tree_seeds, name_map, sample_colors):
    """
    Compute ancestry paths for seed nodes in a tree.

    Args:
        tree: tskit.Tree object
        tree_seeds: List of seed node IDs to trace ancestry
        name_map: Dict mapping node_id to original name
        sample_colors: Optional dict {sample_name: [r,g,b,a]} for coloring

    Returns:
        List of lineage dicts with path_node_ids and color
    """
    tree_lineages = []
    for seed_node in tree_seeds:
        # Trace ancestry path from sample to root
        path_nodes = []
        current = seed_node
        while current != -1 and current != tskit.NULL:
            path_nodes.append(current)
            current = tree.parent(current)

        if len(path_nodes) > 1:
            # Emit root -> tip to match frontend L-shape construction.
            path_nodes = list(reversed(path_nodes))

            # Get color for this lineage
            name = name_map.get(seed_node, str(seed_node))
            color = None
            if sample_colors:
                color = sample_colors.get(name.lower())

            tree_lineages.append({
                "path_node_ids": [int(n) for n in path_nodes],
                "color": color
            })

    return tree_lineages


def search_nodes_in_trees(
    ts,
    sample_names,
    tree_indices,
    show_lineages=False,
    sample_colors=None,
    sample_name_key="name"
):
    """
    Search for nodes matching sample names in specified trees.
    Returns highlights and optionally lineage paths.

    Args:
        ts: tskit.TreeSequence
        sample_names: List of sample names to search for
        tree_indices: List of tree indices to search in
        show_lineages: Whether to compute lineage (ancestry) paths
        sample_colors: Optional dict {sample_name: [r,g,b,a]} for coloring
        sample_name_key: Key in node metadata used as sample name

    Returns:
        dict with:
        - highlights: {tree_idx: [{node_id, name}]}
        - lineage: {tree_idx: [{path: [[x,y]...], color}]} if show_lineages
    """
    if not sample_names or not tree_indices:
        return {"highlights": {}, "lineage": {}}

    # Build sample_name -> node_id mapping
    name_to_node_id = _build_sample_name_mapping(ts, sample_name_key)

    # Convert sample_names to node_ids
    target_node_ids = set()
    name_map = {}  # node_id -> original name
    for name in sample_names:
        lower_name = name.lower()
        if lower_name in name_to_node_id:
            nid = name_to_node_id[lower_name]
            target_node_ids.add(nid)
            name_map[nid] = name

    if not target_node_ids:
        return {"highlights": {}, "lineage": {}}

    highlights = {}
    lineage = {}

    for tree_idx in tree_indices:
        tree_idx = int(tree_idx)
        if tree_idx < 0 or tree_idx >= ts.num_trees:
            continue

        tree = ts.at_index(tree_idx)

        # Find matching samples in this tree
        tree_highlights = []
        tree_seeds = []  # For lineage computation

        for node_id in target_node_ids:
            # Check if this sample is in this tree
            if tree.is_sample(node_id):
                name = name_map.get(node_id, str(node_id))
                tree_highlights.append({
                    "node_id": int(node_id),
                    "name": name
                })
                tree_seeds.append(node_id)

        if tree_highlights:
            highlights[tree_idx] = tree_highlights

            # Compute lineage paths if requested
            if show_lineages and tree_seeds:
                tree_lineages = _compute_lineage_paths(
                    tree, tree_seeds, name_map, sample_colors
                )
                if tree_lineages:
                    lineage[tree_idx] = tree_lineages

    return {"highlights": highlights, "lineage": lineage}


def get_node_details(ts, node_name):
    """Get details for a specific node in the tree sequence."""
    node = ts.node(node_name)
    return {
        "id": node.id,
        "time": node.time,
        "population": node.population,
        "individual": node.individual,
        "metadata": make_json_serializable(node.metadata)
    }


def get_tree_details(ts, tree_index):
    """Get details for a specific tree at the given index."""
    tree = ts.at_index(tree_index)

    mutations = []
    for mut in tree.mutations():
        site = ts.site(mut.site)
        mutations.append({
            "id": mut.id,
            "node": mut.node,  # Node ID for highlighting
            "site_id": mut.site,
            "position": site.position,
            "derived_state": mut.derived_state,
            "inherited_state": ts.mutation(mut.parent).derived_state if mut.parent != -1 else site.ancestral_state
        })

    return {
        "interval": tree.interval,
        "num_roots": tree.num_roots,
        "num_nodes": tree.num_nodes,
        "mutations": mutations
    }


def get_individual_details(ts, individual_id):
    """Get details for a specific individual in the tree sequence."""
    individual = ts.individual(individual_id)
    return {
        "id": individual.id,
        "nodes": make_json_serializable(individual.nodes),
        "metadata": make_json_serializable(individual.metadata)
    }


def get_comprehensive_individual_details(ts, individual_id):
    """Get comprehensive individual table data including location, parents, flags."""
    if individual_id is None or individual_id == -1:
        return None

    individual = ts.individual(individual_id)
    return {
        "id": int(individual.id),
        "flags": int(individual.flags),
        "location": list(individual.location) if len(individual.location) > 0 else None,
        "parents": [int(p) for p in individual.parents] if len(individual.parents) > 0 else [],
        "nodes": [int(n) for n in individual.nodes],
        "metadata": make_json_serializable(individual.metadata)
    }


def get_population_details(ts, population_id):
    """Get population table data."""
    if population_id is None or population_id == -1:
        return None
    pop = ts.population(population_id)
    return {
        "id": int(pop.id),
        "metadata": make_json_serializable(pop.metadata)
    }


def get_mutations_for_node(ts, node_id, tree_index=None):
    """Get all mutations on a specific node, optionally filtered by tree interval."""
    mutations = []

    # Get tree interval if tree_index is specified
    tree_interval = None
    if tree_index is not None:
        tree = ts.at_index(tree_index)
        tree_interval = tree.interval

    for mut in ts.mutations():
        if mut.node == node_id:
            site = ts.site(mut.site)

            # Filter by tree interval if specified
            if tree_interval is not None:
                if not (site.position >= tree_interval.left and site.position < tree_interval.right):
                    continue

            mutations.append({
                "id": int(mut.id),
                "site_id": int(mut.site),
                "position": float(site.position),
                "ancestral_state": site.ancestral_state,
                "derived_state": mut.derived_state,
                "time": float(mut.time) if mut.time != tskit.UNKNOWN_TIME else None,
                "parent_mutation": int(mut.parent) if mut.parent != -1 else None,
                "metadata": make_json_serializable(mut.metadata) if mut.metadata else None
            })

    return mutations


def get_edges_for_node(ts, node_id, tree_index=None):
    """Get all edges where this node is parent or child."""
    edges = {
        "as_parent": [],  # Edges where node is parent
        "as_child": []    # Edges where node is child
    }

    # Get tree interval if tree_index is specified
    tree_interval = None
    if tree_index is not None:
        tree = ts.at_index(tree_index)
        tree_interval = tree.interval

    for edge in ts.edges():
        # Filter by tree interval if specified (edge must overlap with tree)
        if tree_interval is not None:
            if edge.right <= tree_interval.left or edge.left >= tree_interval.right:
                continue

        edge_data = {
            "id": int(edge.id),
            "left": float(edge.left),
            "right": float(edge.right),
            "parent": int(edge.parent),
            "child": int(edge.child)
        }

        if edge.parent == node_id:
            edges["as_parent"].append(edge_data)
        if edge.child == node_id:
            edges["as_child"].append(edge_data)

    return edges


async def handle_details(file_path, data):
    """Handle requests for tree, node, and individual details."""
    try:
        ctx = await get_file_context(file_path)
        if ctx is None:
            return json.dumps({"error": "Tree sequence (ts) is not set. Please upload a file first."})

        ts = ctx.tree_sequence
        return_data = {}
        tree_index = data.get("treeIndex")
        comprehensive = data.get("comprehensive", False)

        if tree_index is not None:
            return_data["tree"] = get_tree_details(ts, int(tree_index))

        node_name = data.get("node")
        if node_name is not None:
            node_id = int(node_name)
            node_details = get_node_details(ts, node_id)
            return_data["node"] = node_details

            # Auto-fetch individual details
            if node_details.get("individual") != -1:
                if comprehensive:
                    return_data["individual"] = get_comprehensive_individual_details(
                        ts, node_details.get("individual")
                    )
                else:
                    return_data["individual"] = get_individual_details(
                        ts, node_details.get("individual")
                    )

            # Comprehensive mode: add population, mutations, edges
            if comprehensive:
                # Population
                if node_details.get("population") != -1:
                    return_data["population"] = get_population_details(
                        ts, node_details.get("population")
                    )

                # Mutations on this node
                return_data["mutations"] = get_mutations_for_node(
                    ts, node_id, tree_index
                )

                # Edges for this node
                return_data["edges"] = get_edges_for_node(
                    ts, node_id, tree_index
                )

        return json.dumps(return_data)
    except Exception as e:
        return json.dumps({"error": f"Error getting details: {str(e)}"})


async def handle_tree_graph_query(
    file_path,
    tree_indices,
    sparsification=False,
    session_id: str = None,
    tree_graph_cache=None,
    csv_tree_graph_cache=None,
    actual_display_array=None
):
    """
    Construct trees using Numba-optimized tree_graph module.

    Args:
        file_path: Path to tree sequence file
        tree_indices: List of tree indices to process
        sparsification: Enable tip-only sparsification (default False)
        session_id: Session ID for cache lookup/storage
        tree_graph_cache: TreeGraphCache instance for caching TreeGraph objects

    Returns:
        dict with:
        - buffer: PyArrow IPC binary data containing:
            - node_id: int32 (tskit node ID)
            - parent_id: int32 (-1 for roots)
            - is_tip: bool
            - tree_idx: int32 (which tree this node belongs to)
            - x: float32 (time-based coordinate [0,1])
            - y: float32 (layout-based coordinate [0,1])
        - global_min_time: float
        - global_max_time: float
        - tree_indices: list[int]
    """
    ctx = await get_file_context(file_path)
    if ctx is None:
        return {"error": "Tree sequence not loaded. Please load a file first."}

    ts = ctx.tree_sequence

    # CSV support: parse Newick strings and build tree layout
    if isinstance(ts, pd.DataFrame):
        shift_tips_to_one = should_shift_csv_tips(ctx.file_path or file_path)
        # Get max_branch_length from config (times.values[1])
        times_values = ctx.config.get("times", {}).get("values", [0.0, 1.0])
        max_branch_length = float(times_values[1]) if len(times_values) > 1 else 1.0
        indices = [int(t) for t in (tree_indices or [])]
        samples_order = ctx.config.get("samples") or []
        pre_parsed_graphs = {}
        if session_id and csv_tree_graph_cache:
            from lorax.csv.newick_tree import parse_newick_to_tree

            for tree_idx in indices:
                cached = await csv_tree_graph_cache.get(session_id, int(tree_idx))
                if cached is not None:
                    pre_parsed_graphs[int(tree_idx)] = cached
                    continue

                # Cache miss: parse and store (best-effort)
                try:
                    newick_str = ts.iloc[int(tree_idx)].get("newick")
                except Exception:
                    newick_str = None
                if newick_str is None or pd.isna(newick_str):
                    continue

                tree_max_branch_length = None
                if "max_branch_length" in ts.columns:
                    try:
                        v = ts.iloc[int(tree_idx)].get("max_branch_length")
                        if v is not None and not (isinstance(v, float) and pd.isna(v)) and str(v).strip() != "":
                            tree_max_branch_length = float(v)
                    except Exception:
                        tree_max_branch_length = None

                try:
                    graph = await asyncio.to_thread(
                        parse_newick_to_tree,
                        str(newick_str),
                        max_branch_length,
                        samples_order=samples_order,
                        tree_max_branch_length=tree_max_branch_length,
                        shift_tips_to_one=shift_tips_to_one,
                    )
                except Exception:
                    continue

                pre_parsed_graphs[int(tree_idx)] = graph
                await csv_tree_graph_cache.set(session_id, int(tree_idx), graph)

            if actual_display_array is not None:
                await csv_tree_graph_cache.evict_not_visible(session_id, set(actual_display_array))

        return build_csv_layout_response(
            ts,
            indices,
            max_branch_length,
            samples_order=samples_order,
            pre_parsed_graphs=pre_parsed_graphs,
            shift_tips_to_one=shift_tips_to_one,
        )

    # Collect pre-cached TreeGraphs
    pre_cached_graphs = {}
    if session_id and tree_graph_cache:
        for tree_idx in tree_indices:
            cached = await tree_graph_cache.get(session_id, int(tree_idx))
            if cached is not None:
                pre_cached_graphs[int(tree_idx)] = cached
        if pre_cached_graphs:
            print(f"TreeGraph cache hits: {len(pre_cached_graphs)}/{len(tree_indices)} trees")

    # Run in thread pool to avoid blocking
    def process_trees():
        return construct_trees_batch(
            ts,
            tree_indices,
            sparsification=sparsification,
            pre_cached_graphs=pre_cached_graphs
        )

    buffer, min_time, max_time, processed_indices, newly_built = await asyncio.to_thread(process_trees)

    # Cache newly built TreeGraphs
    if session_id and tree_graph_cache and newly_built:
        for tree_idx, graph in newly_built.items():
            await tree_graph_cache.set(session_id, tree_idx, graph)
        print(f"TreeGraph cached: {len(newly_built)} new trees for session {session_id[:8]}...")

    # Evict trees no longer in visible set (visibility-based eviction)
    if session_id and tree_graph_cache and actual_display_array is not None:
        await tree_graph_cache.evict_not_visible(session_id, set(actual_display_array))

    return {
        "buffer": buffer,
        "global_min_time": min_time,
        "global_max_time": max_time,
        "tree_indices": processed_indices
    }


async def get_or_construct_tree_graph(
    file_path: str,
    tree_index: int,
    session_id: str,
    tree_graph_cache
) -> TreeGraph:
    """
    Get a TreeGraph from cache or construct and cache it.

    This function is used by lineage operations that need the full TreeGraph
    structure for ancestor/descendant traversal.

    Args:
        file_path: Path to tree sequence file
        tree_index: Index of the tree to get
        session_id: Session ID for cache key
        tree_graph_cache: TreeGraphCache instance

    Returns:
        TreeGraph object, or None if file not loaded
    """
    # Check cache first
    cached = await tree_graph_cache.get(session_id, tree_index)
    if cached is not None:
        print(f"TreeGraph cache hit: session={session_id[:8]}... tree={tree_index}")
        return cached

    # Load file context
    ctx = await get_file_context(file_path)
    if ctx is None:
        return None

    ts = ctx.tree_sequence

    # Can't construct TreeGraph for CSV
    if isinstance(ts, pd.DataFrame):
        return None

    # Construct tree graph
    def _construct():
        edges = ts.tables.edges
        nodes = ts.tables.nodes
        breakpoints = list(ts.breakpoints())
        min_time = float(ts.min_time)
        max_time = float(ts.max_time)
        return construct_tree(ts, edges, nodes, breakpoints, tree_index, min_time, max_time)

    tree_graph = await asyncio.to_thread(_construct)

    # Cache it
    await tree_graph_cache.set(session_id, tree_index, tree_graph)
    print(f"TreeGraph cached: session={session_id[:8]}... tree={tree_index}")

    return tree_graph


async def ensure_trees_cached(
    file_path: str,
    tree_indices: list,
    session_id: str,
    tree_graph_cache
) -> int:
    """
    Ensure multiple trees are cached for a session.

    This is called after process_postorder_layout to cache trees for
    subsequent lineage operations.

    Args:
        file_path: Path to tree sequence file
        tree_indices: List of tree indices to cache
        session_id: Session ID for cache key
        tree_graph_cache: TreeGraphCache instance

    Returns:
        Number of trees newly cached (not already in cache)
    """
    ctx = await get_file_context(file_path)
    if ctx is None:
        return 0

    ts = ctx.tree_sequence

    if isinstance(ts, pd.DataFrame):
        return 0

    newly_cached = 0

    # Pre-extract tables for efficiency
    edges = ts.tables.edges
    nodes = ts.tables.nodes
    breakpoints = list(ts.breakpoints())
    min_time = float(ts.min_time)
    max_time = float(ts.max_time)

    for tree_index in tree_indices:
        tree_index = int(tree_index)

        # Skip if already cached
        cached = await tree_graph_cache.get(session_id, tree_index)
        if cached is not None:
            continue

        # Construct and cache
        def _construct(idx):
            return construct_tree(ts, edges, nodes, breakpoints, idx, min_time, max_time)

        tree_graph = await asyncio.to_thread(_construct, tree_index)
        await tree_graph_cache.set(session_id, tree_index, tree_graph)
        newly_cached += 1

    if newly_cached > 0:
        print(f"Cached {newly_cached} trees for session {session_id[:8]}...")

    return newly_cached


def _get_matching_sample_nodes(ts, metadata_key, metadata_value, sources, sample_name_key):
    """
    Find all sample node IDs that match a metadata value.

    Args:
        ts: tskit.TreeSequence
        metadata_key: Metadata key to filter by
        metadata_value: Metadata value to match
        sources: Metadata sources to search
        sample_name_key: Key in node metadata used as sample name

    Returns:
        Set of matching node IDs
    """
    matching_node_ids = set()
    for node_id in ts.samples():
        sample_name, value = _get_sample_metadata_value(
            ts, node_id, metadata_key, sources, sample_name_key
        )
        if value is not None and str(value) == str(metadata_value):
            matching_node_ids.add(node_id)
    return matching_node_ids


async def _ensure_tree_graph_loaded(
    ts,
    tree_idx,
    session_id,
    tree_graph_cache,
    edges,
    nodes,
    breakpoints,
    min_time,
    max_time
):
    """
    Get tree graph from cache or construct and cache it.

    Args:
        ts: tskit.TreeSequence
        tree_idx: Tree index to load
        session_id: Session ID for cache key
        tree_graph_cache: TreeGraphCache instance
        edges, nodes, breakpoints, min_time, max_time: Pre-extracted table data

    Returns:
        TreeGraph object
    """
    from lorax.tree_graph import construct_tree

    # Try to get from cache first
    graph = await tree_graph_cache.get(session_id, tree_idx)
    if graph is not None:
        return graph

    # Construct tree graph
    def _construct():
        return construct_tree(ts, edges, nodes, breakpoints, tree_idx, min_time, max_time)

    graph = await asyncio.to_thread(_construct)

    # Cache it for future use
    await tree_graph_cache.set(session_id, tree_idx, graph)

    return graph


async def get_highlight_positions(
    ts,
    file_path,
    metadata_key,
    metadata_value,
    tree_indices,
    session_id: str,
    tree_graph_cache,
    sources=("individual", "node", "population"),
    sample_name_key="name"
):
    """
    Get positions for all tip nodes with a specific metadata value.
    Uses cached TreeGraph objects when available.

    Args:
        ts: tskit.TreeSequence
        file_path: Path to tree sequence file (for cache key)
        metadata_key: Metadata key to filter by
        metadata_value: Metadata value to match
        tree_indices: List of tree indices to compute positions for
        session_id: Session ID for cache lookup
        tree_graph_cache: TreeGraphCache instance
        sources: Metadata sources to search
        sample_name_key: Key in node metadata used as sample name

    Returns:
        dict with:
        - positions: List of {node_id, tree_idx, x, y} dicts
    """
    if not tree_indices:
        return {"positions": []}

    # Get sample node IDs that have this metadata value
    matching_node_ids = _get_matching_sample_nodes(
        ts, metadata_key, metadata_value, sources, sample_name_key
    )

    if not matching_node_ids:
        return {"positions": []}

    # Pre-extract tables for reuse (only needed if cache miss)
    edges = ts.tables.edges
    nodes = ts.tables.nodes
    breakpoints = list(ts.breakpoints())
    min_time = float(ts.min_time)
    max_time = float(ts.max_time)

    positions = []

    # For each requested tree, get graph and extract positions
    for tree_idx in tree_indices:
        tree_idx = int(tree_idx)
        if tree_idx < 0 or tree_idx >= ts.num_trees:
            continue

        graph = await _ensure_tree_graph_loaded(
            ts, tree_idx, session_id, tree_graph_cache,
            edges, nodes, breakpoints, min_time, max_time
        )

        # Extract positions for matching nodes that are in this tree
        for node_id in matching_node_ids:
            if graph.in_tree[node_id]:
                positions.append({
                    "node_id": int(node_id),
                    "tree_idx": tree_idx,
                    "x": float(graph.x[node_id]),
                    "y": float(graph.y[node_id])
                })

    return {"positions": positions}


async def get_multi_value_highlight_positions(
    ts,
    file_path,
    metadata_key,
    metadata_values,  # List[str] - Array of values (OR logic)
    tree_indices,
    session_id: str,
    tree_graph_cache,
    show_lineages: bool = False,
    sources=("individual", "node", "population"),
    sample_name_key="name"
):
    """
    Get positions for tip nodes matching ANY of the metadata values.
    Returns positions grouped by value for per-value coloring.

    Args:
        ts: tskit.TreeSequence
        file_path: Path to tree sequence file (for cache key)
        metadata_key: Metadata key to filter by
        metadata_values: List of metadata values to match (OR logic)
        tree_indices: List of tree indices to compute positions for
        session_id: Session ID for cache lookup
        tree_graph_cache: TreeGraphCache instance
        show_lineages: Whether to compute lineage (ancestry) paths
        sources: Metadata sources to search
        sample_name_key: Key in node metadata used as sample name

    Returns:
        dict with:
        - positions_by_value: {"Africa": [{node_id, tree_idx, x, y}, ...], ...}
        - lineages: {"Africa": {tree_idx: [{path_node_ids, color}]}} if show_lineages
        - total_count: int
    """
    if not tree_indices or not metadata_values:
        return {"positions_by_value": {}, "lineages": {}, "total_count": 0}

    # Deduplicate values
    unique_values = list(set(str(v) for v in metadata_values))

    # Pre-extract tables for reuse (only needed if cache miss)
    edges = ts.tables.edges
    nodes = ts.tables.nodes
    breakpoints = list(ts.breakpoints())
    min_time = float(ts.min_time)
    max_time = float(ts.max_time)

    positions_by_value = {}
    lineages = {} if show_lineages else None
    total_count = 0

    # For each value, find matching samples
    for value in unique_values:
        matching_node_ids = _get_matching_sample_nodes(
            ts, metadata_key, value, sources, sample_name_key
        )

        if not matching_node_ids:
            positions_by_value[value] = []
            continue

        value_positions = []
        value_lineages = {} if show_lineages else None

        # For each requested tree, get graph and extract positions
        for tree_idx in tree_indices:
            tree_idx = int(tree_idx)
            if tree_idx < 0 or tree_idx >= ts.num_trees:
                continue

            graph = await _ensure_tree_graph_loaded(
                ts, tree_idx, session_id, tree_graph_cache,
                edges, nodes, breakpoints, min_time, max_time
            )

            tree_positions = []
            tree_seeds = []  # For lineage computation

            # Extract positions for matching nodes that are in this tree
            for node_id in matching_node_ids:
                if graph.in_tree[node_id]:
                    tree_positions.append({
                        "node_id": int(node_id),
                        "tree_idx": tree_idx,
                        "x": float(graph.x[node_id]),
                        "y": float(graph.y[node_id])
                    })
                    tree_seeds.append(node_id)

            value_positions.extend(tree_positions)

            # Compute lineage paths if requested
            if show_lineages and tree_seeds:
                tree = ts.at_index(tree_idx)
                name_map = {nid: str(nid) for nid in tree_seeds}
                tree_lineages = _compute_lineage_paths(
                    tree, tree_seeds, name_map, None  # No per-sample colors, use value color
                )
                if tree_lineages:
                    value_lineages[tree_idx] = tree_lineages

        positions_by_value[value] = value_positions
        total_count += len(value_positions)

        if show_lineages and value_lineages:
            lineages[value] = value_lineages

    result = {
        "positions_by_value": positions_by_value,
        "total_count": total_count
    }

    if show_lineages:
        result["lineages"] = lineages

    return result
