"""
LiquidFunctionClient - Client for liquid template analysis and operations.

This client provides functions for:
1. Building liquid template structure documentation
2. Analyzing bad cases to identify priority templates for optimization
3. Managing template structure caching
"""
from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import os
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from dotenv import load_dotenv

from prompt_auto_lab.core.store_dir import get_store_dir

load_dotenv()

module_logger = logging.getLogger(__name__)


def _parse_conditional_blocks(source: str) -> list[dict]:
    """Parse Liquid conditional blocks from source, returning structured block info.

    Returns list of dicts with keys:
        block_start: 1-based line number of opening tag ({% if %} or {% unless %})
        block_end: 1-based line number of closing tag ({% endif %} or {% endunless %})
        branches: list of dicts with keys:
            type: 'if' | 'elsif' | 'else' | 'unless'
            start: 1-based line number of the branch content start
            end: 1-based line number of the branch content end (inclusive)
    """
    lines = source.split("\n")
    tag_re = re.compile(
        r'^\s*\{%-?\s*(if|elsif|else|unless|endif|endunless)\b.*?-?%\}\s*$'
    )

    stack: list[dict] = []
    blocks: list[dict] = []

    for i, line in enumerate(lines):
        line_num = i + 1  # 1-based
        m = tag_re.match(line)
        if not m:
            continue
        tag = m.group(1)

        if tag in ("if", "unless"):
            stack.append({
                "block_start": line_num,
                "tag_type": tag,
                "branches": [{"type": tag, "start": line_num + 1, "end": None}],
            })
        elif tag == "elsif" and stack:
            stack[-1]["branches"][-1]["end"] = line_num - 1
            stack[-1]["branches"].append(
                {"type": "elsif", "start": line_num + 1, "end": None}
            )
        elif tag == "else" and stack:
            stack[-1]["branches"][-1]["end"] = line_num - 1
            stack[-1]["branches"].append(
                {"type": "else", "start": line_num + 1, "end": None}
            )
        elif tag in ("endif", "endunless") and stack:
            stack[-1]["branches"][-1]["end"] = line_num - 1
            block = stack.pop()
            block["block_end"] = line_num
            blocks.append(block)

    return blocks


def build_filtered_content(
    source: str,
    branch_mapping: list[dict],
) -> dict:
    """Build filtered Liquid content keeping only active conditional branches.

    Args:
        source: Raw Liquid template source text.
        branch_mapping: List of dicts from LLM branch mapping, each with:
            line_start: 1-based start line of the conditional block
            line_end: 1-based end line of the conditional block
            active_branch: 'if' | 'elsif' | 'else' | 'unless' -- which branch is active
            confidence: 'high' | 'low' | 'unknown'

    Returns:
        Dict with:
            filtered_content: Filtered source with inactive branches removed.
            active_line_ranges: List of dicts with line_start/line_end (1-based)
                for each active branch's content region in the original source.
            line_map: List of 1-based original line numbers. line_map[i] is the
                original source line that corresponds to filtered output line i+1.
                Empty when no filtering was applied.
    """
    if not branch_mapping:
        return {"filtered_content": source, "active_line_ranges": [], "line_map": []}

    # Index mappings by (line_start, line_end) for lookup
    mapping_by_range: dict[tuple[int, int], dict] = {}
    for m in branch_mapping:
        key = (m["line_start"], m["line_end"])
        mapping_by_range[key] = m

    # Parse actual conditional blocks from source
    blocks = _parse_conditional_blocks(source)

    # Determine which lines to remove
    lines = source.split("\n")
    lines_to_remove: set[int] = set()  # 1-based line numbers
    active_line_ranges: list[dict] = []

    for block in reversed(blocks):
        key = (block["block_start"], block["block_end"])
        mapping = mapping_by_range.get(key)

        if not mapping or mapping.get("confidence", "low") != "high":
            continue  # Preserve entire block as-is

        active_type = mapping["active_branch"]

        # __none__ means the entire block was inactive — remove all lines
        if active_type == "__none__":
            for ln in range(block["block_start"], block["block_end"] + 1):
                lines_to_remove.add(ln)
            continue

        # Find the active branch
        active_branch = None
        for branch in block["branches"]:
            if branch["type"] == active_type:
                active_branch = branch
                break

        if active_branch is None:
            continue  # Can't find active branch, preserve as-is

        # Remove conditional tag lines
        lines_to_remove.add(block["block_start"])
        lines_to_remove.add(block["block_end"])

        # Remove inactive branch content + their tag lines
        for branch in block["branches"]:
            if branch["type"] == active_type:
                continue  # Keep active branch content
            # Remove the tag line for this branch (elsif/else)
            tag_line = branch["start"] - 1
            if tag_line > block["block_start"]:
                lines_to_remove.add(tag_line)
            # Remove content lines
            for ln in range(
                branch["start"], (branch["end"] or branch["start"]) + 1
            ):
                lines_to_remove.add(ln)

        # Remove the active branch's own tag line (elsif/else tag leaks otherwise)
        if active_branch and active_type in ("elsif", "else"):
            active_tag_line = active_branch["start"] - 1
            if active_tag_line > block["block_start"]:
                lines_to_remove.add(active_tag_line)

        # Record active branch line range
        active_line_ranges.append({
            "line_start": active_branch["start"],
            "line_end": active_branch["end"] or active_branch["start"],
            "block_start": block["block_start"],
            "block_end": block["block_end"],
        })

    # Build result keeping only non-removed lines, and build line_map
    result_lines = []
    line_map = []  # line_map[i] = 1-based original line number for filtered line i+1
    for i, line in enumerate(lines):
        if (i + 1) not in lines_to_remove:
            result_lines.append(line)
            line_map.append(i + 1)

    return {
        "filtered_content": "\n".join(result_lines),
        "active_line_ranges": active_line_ranges,
        "line_map": line_map,
    }


def _normalize_for_matching(text: str) -> str:
    """Normalize text for substring matching by stripping Liquid tags and collapsing whitespace.

    Removes:
        - Block tags: {% if %}, {% endif %}, {% include %}, {% assign %}, etc.
        - Output tags: {{ variable }}
        - Leading/trailing whitespace
        - Collapses multiple whitespace (including newlines) to single space
    """
    # Remove Liquid block tags: {%- ... -%} or {% ... %}
    result = re.sub(r'\{%-?.*?-?%\}', '', text)
    # Remove Liquid output tags: {{ ... }}
    result = re.sub(r'\{\{.*?\}\}', '', result)
    # Collapse all whitespace (spaces, tabs, newlines) to single space
    result = re.sub(r'\s+', ' ', result).strip()
    return result


def _extract_branch_texts(
    source: str,
    block: dict,
) -> list[tuple[str, str]]:
    """Extract normalized plain text for each branch in a conditional block.

    Args:
        source: Full Liquid template source text.
        block: A block dict from _parse_conditional_blocks() with
            'branches' list of {'type', 'start', 'end'}.

    Returns:
        List of (branch_type, normalized_text) tuples.
        branch_type is 'if', 'elsif', 'else', or 'unless'.
        normalized_text has Liquid tags stripped and whitespace collapsed.
    """
    lines = source.split("\n")
    result = []
    for branch in block["branches"]:
        start = branch["start"]  # 1-based, first content line
        end = branch["end"] or start  # 1-based, last content line (inclusive)
        branch_lines = lines[start - 1 : end]
        raw_text = "\n".join(branch_lines)
        normalized = _normalize_for_matching(raw_text)
        result.append((branch["type"], normalized))
    return result


def match_conditional_branches(
    source: str,
    rendered_prompt: str,
) -> list[dict]:
    """Deterministically match conditional branches against rendered output.

    For each conditional block in the source, extracts plain text from each
    branch and checks which branch's text appears in the rendered prompt.

    Handles nested blocks by extracting only the "own" lines of each branch
    (excluding nested conditional block ranges) for matching.

    Args:
        source: Raw Liquid template source text.
        rendered_prompt: Fully rendered prompt text from evaluation.

    Returns:
        List of dicts compatible with build_filtered_content() branch_mapping:
            line_start: 1-based start line of the conditional block
            line_end: 1-based end line of the conditional block
            active_branch: 'if' | 'elsif' | 'else' | 'unless'
            confidence: 'high' | 'low'
            match_method: 'string_match' (for debugging)
    """
    blocks = _parse_conditional_blocks(source)
    if not blocks:
        return []

    # Build a set of nested block ranges for each outer block so we can
    # exclude nested content when matching outer branches.
    nested_ranges_by_block: dict[tuple[int, int], list[tuple[int, int]]] = {}
    for block in blocks:
        key = (block["block_start"], block["block_end"])
        nested = []
        for other in blocks:
            if other is block:
                continue
            # 'other' is nested inside 'block' if fully contained
            if other["block_start"] > block["block_start"] and other["block_end"] < block["block_end"]:
                nested.append((other["block_start"], other["block_end"]))
        nested_ranges_by_block[key] = nested

    lines = source.split("\n")
    normalized_rendered = _normalize_for_matching(rendered_prompt)
    results = []

    for block in blocks:
        key = (block["block_start"], block["block_end"])
        nested_ranges = nested_ranges_by_block[key]

        if not nested_ranges:
            # No nesting -- use standard _extract_branch_texts
            branch_texts = _extract_branch_texts(source, block)
        else:
            # Extract branch texts excluding nested block line ranges
            branch_texts = _extract_branch_texts_excluding(
                lines, block, nested_ranges
            )

        # Find which branches have their text in the rendered output
        matches = []
        for branch_type, normalized_text in branch_texts:
            if not normalized_text:
                continue  # Skip empty branches
            if normalized_text in normalized_rendered:
                matches.append(branch_type)

        if len(matches) == 1:
            confidence = "high"
            active_branch = matches[0]
            module_logger.debug(
                "Block lines %d-%d: matched branch '%s' (text found in rendered output)",
                block["block_start"], block["block_end"], active_branch,
            )
        else:
            # 0 matches (variables substituted) or 2+ matches (ambiguous)
            confidence = "low"
            active_branch = block["branches"][0]["type"]  # default to first
            module_logger.debug(
                "Block lines %d-%d: %d branches matched (%s), confidence=low",
                block["block_start"], block["block_end"],
                len(matches), matches or "none",
            )

        results.append({
            "line_start": block["block_start"],
            "line_end": block["block_end"],
            "active_branch": active_branch,
            "confidence": confidence,
            "match_method": "string_match",
        })

    return results


def _extract_branch_texts_excluding(
    lines: list[str],
    block: dict,
    nested_ranges: list[tuple[int, int]],
) -> list[tuple[str, str]]:
    """Extract normalized branch texts, excluding lines inside nested blocks.

    Args:
        lines: Source split into lines (0-indexed list).
        block: A block dict from _parse_conditional_blocks().
        nested_ranges: List of (block_start, block_end) 1-based line ranges
            for nested conditional blocks to exclude.

    Returns:
        List of (branch_type, normalized_text) tuples, same as _extract_branch_texts.
    """
    # Build set of 1-based line numbers to exclude
    excluded: set[int] = set()
    for rng_start, rng_end in nested_ranges:
        for ln in range(rng_start, rng_end + 1):
            excluded.add(ln)

    result = []
    for branch in block["branches"]:
        start = branch["start"]  # 1-based
        end = branch["end"] or start  # 1-based inclusive
        own_lines = []
        for ln in range(start, end + 1):
            if ln not in excluded:
                own_lines.append(lines[ln - 1])
        raw_text = "\n".join(own_lines)
        normalized = _normalize_for_matching(raw_text)
        result.append((branch["type"], normalized))
    return result


def _build_block_tree(blocks: list[dict]) -> list[dict]:
    """Convert flat block list into tree based on line-range containment.

    Each node gets added fields:
        children: list[dict]  — nested blocks in source order
        parent_branch_type: str | None  — which branch of parent contains this block
    """
    # Sort by block_start so parents come before children
    sorted_blocks = sorted(blocks, key=lambda b: b["block_start"])

    # Add tree fields to each block
    for block in sorted_blocks:
        block["children"] = []
        block["parent_branch_type"] = None

    roots: list[dict] = []

    for block in sorted_blocks:
        # Find the tightest parent (smallest range that fully contains this block)
        best_parent = None
        for candidate in sorted_blocks:
            if candidate is block:
                continue
            if (candidate["block_start"] < block["block_start"]
                    and candidate["block_end"] > block["block_end"]):
                if best_parent is None or (
                    candidate["block_end"] - candidate["block_start"]
                    < best_parent["block_end"] - best_parent["block_start"]
                ):
                    best_parent = candidate

        if best_parent is None:
            roots.append(block)
        else:
            best_parent["children"].append(block)
            # Determine which branch of parent contains this block
            for branch in best_parent["branches"]:
                branch_start = branch["start"]
                branch_end = branch["end"] or branch_start
                if block["block_start"] >= branch_start and block["block_end"] <= branch_end:
                    block["parent_branch_type"] = branch["type"]
                    break

    return roots


def _extract_branch_fragments(
    lines: list[str],
    branch: dict,
    excluded_line_ranges: list[tuple[int, int]] | None = None,
) -> list[str]:
    """Extract ordered static text fragments from a branch.

    Steps:
    1. Collect branch content lines (excluding nested block ranges)
    2. For each line: remove {% block_tags %}, split on {{ variable }}
    3. Normalize each fragment (collapse whitespace)
    4. Filter out fragments shorter than 3 chars
    """
    start = branch["start"]  # 1-based
    end = branch["end"] or start  # 1-based inclusive

    # Build set of excluded line numbers
    excluded: set[int] = set()
    if excluded_line_ranges:
        for rng_start, rng_end in excluded_line_ranges:
            for ln in range(rng_start, rng_end + 1):
                excluded.add(ln)

    # Process each line individually to produce per-line fragments
    fragments = []
    for ln in range(start, end + 1):
        if ln in excluded:
            continue
        raw_line = lines[ln - 1]

        # Remove block tags: {%- ... -%} or {% ... %}
        cleaned = re.sub(r'\{%-?.*?-?%\}', '', raw_line)

        # Split on output tags: {{ ... }}
        parts = re.split(r'\{\{.*?\}\}', cleaned)

        for part in parts:
            normalized = re.sub(r'\s+', ' ', part).strip()
            if len(normalized) >= 3:
                fragments.append(normalized)

    return fragments


def _match_fragments_at_cursor(
    normalized_rendered: str,
    cursor: int,
    fragments: list[str],
    min_match_ratio: float = 0.6,
) -> tuple[bool, int, int, int]:
    """Check if fragments appear in order at/after cursor position.

    Uses best-effort matching: skips fragments that aren't found and counts
    how many matched. Returns success if at least min_match_ratio of fragments
    matched (minimum 1). Forward-only cursor prevents false positives.

    Returns (matched, new_cursor_position, matched_count, total_count).
    """
    if not fragments:
        return False, cursor, 0, 0

    matched_count = 0
    pos = cursor
    last_matched_pos = cursor

    for fragment in fragments:
        idx = normalized_rendered.find(fragment, pos)
        if idx >= 0:
            matched_count += 1
            pos = idx + len(fragment)
            last_matched_pos = pos

    total = len(fragments)
    threshold = max(1, int(total * min_match_ratio))
    if matched_count >= threshold:
        return True, last_matched_pos, matched_count, total

    return False, cursor, matched_count, total


def _process_block_recursive(
    block: dict,
    lines: list[str],
    normalized_rendered: str,
    cursor: int,
    parent_active: bool,
    results: list[dict],
) -> int:
    """Process a block and its children recursively.

    Args:
        block: Block dict with children from _build_block_tree().
        lines: Source lines (0-indexed list).
        normalized_rendered: Normalized rendered prompt text.
        cursor: Current cursor position in normalized_rendered.
        parent_active: Whether this block's parent branch is active.
        results: List to append match results to (mutated in place).

    Returns:
        Updated cursor position.
    """
    # Compute nested line ranges for this block (direct children only)
    child_ranges = [
        (child["block_start"], child["block_end"]) for child in block["children"]
    ]

    if not parent_active:
        # Parent is inactive → entire block is inactive
        results.append({
            "line_start": block["block_start"],
            "line_end": block["block_end"],
            "active_branch": "__none__",
            "confidence": "high",
            "match_method": "inactive_parent",
        })
        # Recursively mark all children as inactive too
        for child in block["children"]:
            cursor = _process_block_recursive(
                child, lines, normalized_rendered, cursor,
                parent_active=False, results=results,
            )
        return cursor

    # Extract fragments for each branch (excluding nested block ranges)
    branch_fragments: list[tuple[str, list[str]]] = []
    for branch in block["branches"]:
        # Only exclude children that belong to this branch
        branch_child_ranges = [
            (child["block_start"], child["block_end"])
            for child in block["children"]
            if child.get("parent_branch_type") == branch["type"]
        ]
        fragments = _extract_branch_fragments(lines, branch, branch_child_ranges)
        branch_fragments.append((branch["type"], fragments))

    # Single-branch block ({% if %} without else) — check if absent
    if len(block["branches"]) == 1:
        branch_type, fragments = branch_fragments[0]
        if fragments:
            matched, new_cursor, mc, tc = _match_fragments_at_cursor(
                normalized_rendered, cursor, fragments,
            )
            if not matched:
                # Content not found → condition was false
                results.append({
                    "line_start": block["block_start"],
                    "line_end": block["block_end"],
                    "active_branch": "__none__",
                    "confidence": "high",
                    "match_method": "single_branch_absent",
                })
                # Children are inside the inactive branch
                for child in block["children"]:
                    cursor = _process_block_recursive(
                        child, lines, normalized_rendered, cursor,
                        parent_active=False, results=results,
                    )
                return cursor
            else:
                # Content found → condition was true
                results.append({
                    "line_start": block["block_start"],
                    "line_end": block["block_end"],
                    "active_branch": branch_type,
                    "confidence": "high",
                    "match_method": "sequential_match",
                })
                cursor = new_cursor
                # Process children (they are active)
                for child in block["children"]:
                    cursor = _process_block_recursive(
                        child, lines, normalized_rendered, cursor,
                        parent_active=True, results=results,
                    )
                return cursor
        else:
            # No static fragments to match — low confidence
            results.append({
                "line_start": block["block_start"],
                "line_end": block["block_end"],
                "active_branch": branch_type,
                "confidence": "low",
                "match_method": "sequential_match",
            })
            for child in block["children"]:
                cursor = _process_block_recursive(
                    child, lines, normalized_rendered, cursor,
                    parent_active=True, results=results,
                )
            return cursor

    # Multi-branch block — try each branch's fragments, pick best match
    # Save cursor before fragment matching.  Parent fragments come from
    # lines *outside* child ranges.  In the rendered output these lines
    # often appear **after** child content (source order), so the cursor
    # after parent matching may have advanced past children's text region.
    # Children must therefore start from the pre-match position.
    cursor_before_match = cursor

    # Collect (branch_type, new_cursor, matched_count, total_count) for each
    match_results: list[tuple[str, int, int, int]] = []
    for branch_type, fragments in branch_fragments:
        if not fragments:
            match_results.append((branch_type, cursor, 0, 0))
            continue
        matched, new_cursor, mc, tc = _match_fragments_at_cursor(
            normalized_rendered, cursor, fragments,
        )
        if matched:
            match_results.append((branch_type, new_cursor, mc, tc))
        else:
            match_results.append((branch_type, cursor, mc, tc))

    # Find branches that passed the threshold
    passed = [(bt, nc, mc, tc) for bt, nc, mc, tc in match_results if mc >= max(1, int(tc * 0.6)) and tc > 0]

    # Check if any branch had 0 fragments (all content in children).
    # When that happens, we cannot reliably distinguish branches by fragment
    # matching alone because the branch with fragments may match on shared
    # text (e.g. headers that appear in both branches).  Treat as ambiguous.
    has_empty_branch = any(tc == 0 for _, _, _, tc in match_results)

    if has_empty_branch:
        # A branch had 0 fragments (all its content is in children).
        # We cannot reliably distinguish branches by fragment matching
        # because the branch with fragments may match on shared text
        # (e.g. headers identical in both branches).  Mark low confidence
        # and let all children be processed as active.
        #
        # Exception: if the empty-fragment branch is truly empty (has no
        # children either), it's just a blank branch (e.g. {% if x %}{% else
        # %}content{% endif %}) and the non-empty branch is a safe pick.
        truly_empty = True
        for bt, _, mc, tc in match_results:
            if tc == 0:
                # Check whether this branch has children
                branch_has_children = any(
                    child.get("parent_branch_type") == bt
                    for child in block["children"]
                )
                if branch_has_children:
                    truly_empty = False
                    break

        if not truly_empty:
            results.append({
                "line_start": block["block_start"],
                "line_end": block["block_end"],
                "active_branch": block["branches"][0]["type"],
                "confidence": "low",
                "match_method": "content_in_children",
                "branches": [
                    {
                        "type": bt,
                        "start": next(b["start"] for b in block["branches"] if b["type"] == bt),
                        "end": next(b.get("end") for b in block["branches"] if b["type"] == bt),
                        "matched_count": mc,
                        "total_fragments": tc,
                    }
                    for bt, _, mc, tc in match_results
                ],
            })
            for child in block["children"]:
                cursor = _process_block_recursive(
                    child, lines, normalized_rendered, cursor,
                    parent_active=True, results=results,
                )
            return cursor
        # else: truly empty branch — fall through to normal matching below

    if len(passed) == 1:
        active_type, new_cursor, _, _ = passed[0]
        results.append({
            "line_start": block["block_start"],
            "line_end": block["block_end"],
            "active_branch": active_type,
            "confidence": "high",
            "match_method": "sequential_match",
        })
        child_cursor = cursor_before_match
        for child in block["children"]:
            child_active = child.get("parent_branch_type") == active_type
            child_cursor = _process_block_recursive(
                child, lines, normalized_rendered, child_cursor,
                parent_active=child_active, results=results,
            )
        cursor = max(new_cursor, child_cursor)
    elif len(passed) == 0:
        # No branches matched enough
        all_empty = all(tc == 0 for _, _, _, tc in match_results)
        results.append({
            "line_start": block["block_start"],
            "line_end": block["block_end"],
            "active_branch": block["branches"][0]["type"],
            "confidence": "low",
            "match_method": "sequential_match",
        })
        for child in block["children"]:
            cursor = _process_block_recursive(
                child, lines, normalized_rendered, cursor,
                parent_active=True, results=results,
            )
    else:
        # Multiple branches passed threshold — pick highest match ratio
        best = max(passed, key=lambda x: x[2] / x[3] if x[3] > 0 else 0)
        best_ratio = best[2] / best[3] if best[3] > 0 else 0
        # Check if the best is clearly better than second-best
        second_best_ratio = 0
        for bt, nc, mc, tc in passed:
            if bt != best[0]:
                ratio = mc / tc if tc > 0 else 0
                second_best_ratio = max(second_best_ratio, ratio)

        if best_ratio - second_best_ratio >= 0.1:
            # Clear winner
            active_type, new_cursor = best[0], best[1]
            results.append({
                "line_start": block["block_start"],
                "line_end": block["block_end"],
                "active_branch": active_type,
                "confidence": "high",
                "match_method": "sequential_match",
            })
            child_cursor = cursor_before_match
            for child in block["children"]:
                child_active = child.get("parent_branch_type") == active_type
                child_cursor = _process_block_recursive(
                    child, lines, normalized_rendered, child_cursor,
                    parent_active=child_active, results=results,
                )
            cursor = max(new_cursor, child_cursor)
        else:
            # Ambiguous — low confidence
            results.append({
                "line_start": block["block_start"],
                "line_end": block["block_end"],
                "active_branch": block["branches"][0]["type"],
                "confidence": "low",
                "match_method": "sequential_match",
            })
            for child in block["children"]:
                cursor = _process_block_recursive(
                    child, lines, normalized_rendered, cursor,
                    parent_active=True, results=results,
                )

    return cursor


def _mark_descendants_inactive(block: dict, result_index: dict[tuple[int, int], dict]) -> None:
    """Recursively mark all descendants of a block as inactive in result_index."""
    for child in block.get("children", []):
        key = (child["block_start"], child["block_end"])
        cr = result_index.get(key)
        if cr:
            cr["active_branch"] = "__none__"
            cr["confidence"] = "high"
            cr["match_method"] = "inactive_parent"
        _mark_descendants_inactive(child, result_index)


def _infer_from_children(
    roots: list[dict],
    results: list[dict],
) -> None:
    """Post-process results to infer parent branch activity from child results.

    When a parent block was marked low-confidence due to content-in-children
    (one branch has 0 own fragments but has children), we can infer which
    parent branch is active by examining child results:
    - If children in one branch have high-confidence active matches and
      children in other branches are all __none__ or absent → upgrade parent.

    Mutates results in place.
    """
    result_index: dict[tuple[int, int], dict] = {}
    for r in results:
        result_index[(r["line_start"], r["line_end"])] = r

    def _walk(blocks: list[dict]) -> None:
        """Walk tree bottom-up, upgrade content_in_children blocks."""
        for block in blocks:
            # Recurse into children first (bottom-up) so nested inferences
            # resolve before parent checks.
            _walk(block.get("children", []))

            key = (block["block_start"], block["block_end"])
            parent_result = result_index.get(key)

            # Only process content-in-children candidates
            if parent_result and parent_result.get("match_method") == "content_in_children":
                # Group direct children by parent_branch_type
                children_by_branch: dict[str, list[dict]] = {}
                for child in block.get("children", []):
                    bt = child.get("parent_branch_type", "")
                    children_by_branch.setdefault(bt, []).append(child)

                # Classify each branch's children
                branch_has_active: dict[str, bool] = {}
                for bt, children in children_by_branch.items():
                    has_active = False
                    for child in children:
                        cr = result_index.get((child["block_start"], child["block_end"]))
                        if cr and cr.get("confidence") == "high" and cr.get("active_branch") != "__none__":
                            has_active = True
                            break
                    branch_has_active[bt] = has_active

                # Infer: exactly one branch has active children,
                # all other branches with children have no active children
                active_branches = [bt for bt, has in branch_has_active.items() if has]
                if len(active_branches) == 1:
                    inferred = active_branches[0]
                    others_with_children = [
                        bt for bt in branch_has_active
                        if bt != inferred and len(children_by_branch[bt]) > 0
                    ]
                    if all(not branch_has_active[bt] for bt in others_with_children):
                        # Upgrade parent
                        parent_result["active_branch"] = inferred
                        parent_result["confidence"] = "high"
                        parent_result["match_method"] = "child_inference"

                        # Mark children in non-active branches as inactive
                        for bt in others_with_children:
                            for child in children_by_branch[bt]:
                                cr = result_index.get((child["block_start"], child["block_end"]))
                                if cr:
                                    cr["active_branch"] = "__none__"
                                    cr["confidence"] = "high"
                                    cr["match_method"] = "inactive_parent"
                                _mark_descendants_inactive(child, result_index)

    _walk(roots)


def match_conditional_branches_sequential(
    source: str,
    rendered_prompt: str,
) -> list[dict]:
    """Sequential cursor-based matching of conditional branches against rendered output.

    Exploits the fact that rendered prompt content appears in source order.
    Builds a block tree and processes recursively, propagating inactive
    parent status to nested blocks.

    Args:
        source: Raw Liquid template source text.
        rendered_prompt: Fully rendered prompt text from evaluation.

    Returns:
        List of dicts compatible with build_filtered_content() branch_mapping:
            line_start: 1-based start line of the conditional block
            line_end: 1-based end line of the conditional block
            active_branch: 'if' | 'elsif' | 'else' | 'unless' | '__none__'
            confidence: 'high' | 'low'
            match_method: 'sequential_match' | 'inactive_parent' | 'single_branch_absent'
    """
    blocks = _parse_conditional_blocks(source)
    if not blocks:
        return []

    lines = source.split("\n")

    # Normalize rendered prompt (strip Liquid tags, collapse whitespace)
    normalized_rendered = _normalize_for_matching(rendered_prompt)

    # Build block tree
    roots = _build_block_tree(blocks)

    results: list[dict] = []
    cursor = 0

    # Process root blocks in source order
    for root in roots:
        cursor = _process_block_recursive(
            root, lines, normalized_rendered, cursor,
            parent_active=True, results=results,
        )

    # Post-processing: infer parent branch from child results
    _infer_from_children(roots, results)

    return results


class LiquidFunctionClient:
    """
    Client for liquid template analysis operations.

    Provides functions for:
    - Building include trees from entry point templates
    - Generating structure documentation
    - Analyzing bad cases to identify priority templates
    - Caching and reusing template analysis
    """

    def __init__(
        self,
        repo_path: str | None = None,
        store_dir: str | None = None,
    ):
        """
        Initialize the LiquidFunctionClient.

        Args:
            repo_path: Path to the Picasso repository
            store_dir: Directory for storing output files (run_id specific)
        """
        # Picasso repository path
        if repo_path is None:
            raise ValueError("repo_path is required for LiquidFunctionClient")
        self.repo_path = str(repo_path)

        # Store directory
        if store_dir is None:
            store_dir = str(get_store_dir())
        self.store_dir = str(store_dir)

        os.makedirs(self.store_dir, exist_ok=True)

        # Templates directory in Picasso
        self.templates_base = os.path.join(self.repo_path, "Service.Chat", "Inference", "Prompts", "Templates")

        # Structure doc path (use Path object for easier manipulation)
        self._structure_doc_path = Path(self.store_dir) / "liquid_structure.md"

        # Knowledge base for caching
        self._knowledge_base_path = os.path.join(self.store_dir, "template_knowledge.json")
        self._knowledge_base: dict[str, Any] | None = None

        # Cache for condition-aware filtering (populated by build_and_write_structure_doc*)
        self._last_all_templates: list[dict[str, Any]] = []

    @property
    def structure_doc_path(self) -> str:
        """Get the path to the structure documentation file."""
        return str(self._structure_doc_path)

    def structure_doc_exists(self) -> bool:
        """Check if the structure documentation already exists."""
        return os.path.exists(self._structure_doc_path)

    def read_structure_doc(self) -> str | None:
        """Read the existing structure documentation if it exists."""
        if self.structure_doc_exists():
            with open(self._structure_doc_path, "r", encoding="utf-8") as f:
                return f.read()
        return None

    def _derive_entry_base(self, entry_path: str) -> str | None:
        """
        Derive the entry-scoped resolution base from an entry path.

        If the path contains '/LiquidTemplates/', uses everything up to and
        including 'LiquidTemplates/' as the base (repo-absolute).
        Otherwise returns None (fall back to self.templates_base).

        Args:
            entry_path: Entry point path (relative to repo or absolute)

        Returns:
            Absolute path to the entry base directory, or None
        """
        # Normalize to forward slashes for matching
        normalized = entry_path.replace("\\", "/")
        marker = "/LiquidTemplates/"
        idx = normalized.find(marker)
        if idx == -1:
            # Also check at start (path begins with LiquidTemplates/)
            if normalized.startswith("LiquidTemplates/"):
                rel = "LiquidTemplates/"
                base = os.path.join(self.repo_path, rel)
                if os.path.isdir(base):
                    return base
                module_logger.warning(
                    "Entry base directory does not exist: %s (falling back to templates_base)", base
                )
            return None
        # Everything up to and including LiquidTemplates/
        rel = normalized[: idx + len(marker)]
        base = os.path.join(self.repo_path, rel)
        if os.path.isdir(base):
            return base
        module_logger.warning(
            "Entry base directory does not exist: %s (falling back to templates_base)", base
        )
        return None

    def _resolve_template_path(self, include_path: str, current_file: str, entry_base: str | None = None) -> str | None:
        """
        Resolve an include path to an absolute file path.

        Liquid includes can be:
        - Relative to entry base (for tool prompts under LiquidTemplates/)
        - Relative to templates base: 'responding-common/clarifying'
        - Relative to current file: './partial'
        - With or without .liquid extension

        Args:
            include_path: The include path from the liquid template
            current_file: Path to the file containing the include directive
            entry_base: Optional entry-scoped resolution base (derived from LiquidTemplates/ path)
        """
        # Remove quotes if present
        include_path = include_path.strip("'\"")

        # Add .liquid extension if not present
        if not include_path.endswith(".liquid"):
            include_path = include_path + ".liquid"

        # Try different resolution strategies
        current_dir = os.path.dirname(current_file)
        parent_dir = os.path.dirname(current_dir)

        candidates = []

        if entry_base:
            # Entry-scoped base first (for tool prompts)
            candidates.append(os.path.join(entry_base, include_path))
            candidates.extend([
                # Relative to current file's directory
                os.path.join(current_dir, include_path),
                # Relative to current file's parent directory
                os.path.join(parent_dir, include_path),
                # Relative to global templates base (fallback)
                os.path.join(self.templates_base, include_path),
            ])
        else:
            # Legacy resolution order: templates_base first (backward compatible)
            candidates.extend([
                os.path.join(self.templates_base, include_path),
                os.path.join(current_dir, include_path),
                os.path.join(parent_dir, include_path),
            ])

        for candidate in candidates:
            if os.path.exists(candidate):
                return os.path.abspath(candidate)

        return None

    def _read_liquid_file(self, file_path: str, entry_base: str | None = None) -> dict[str, Any]:
        """
        Read content of a Liquid template file.

        Args:
            file_path: Path to the liquid file (can be relative to entry base/templates base, or absolute)
            entry_base: Optional entry-scoped resolution base (tried before templates_base)

        Returns:
            {"ok": bool, "content": str, "path": str, "size": int} or {"ok": False, "error": str}
        """
        try:
            # If not absolute, try candidate paths in order
            if not os.path.isabs(file_path):
                if not file_path.endswith(".liquid"):
                    file_path = file_path + ".liquid"
                if entry_base:
                    # entry_base (for include-relative), repo_path (for repo-relative), templates_base (legacy)
                    candidates = [
                        os.path.join(entry_base, file_path),
                        os.path.join(self.repo_path, file_path),
                        os.path.join(self.templates_base, file_path),
                    ]
                else:
                    # repo_path (for full repo-relative paths), templates_base (for legacy short paths)
                    candidates = [
                        os.path.join(self.repo_path, file_path),
                        os.path.join(self.templates_base, file_path),
                    ]
                resolved = None
                for candidate in candidates:
                    if os.path.exists(candidate):
                        resolved = candidate
                        break
                file_path = resolved if resolved else candidates[0]

            if not os.path.exists(file_path):
                return {"ok": False, "error": f"File not found: {file_path}"}

            with open(file_path, "r", encoding="utf-8") as f:
                content = f.read()

            # Calculate relative path
            abs_path = os.path.abspath(file_path)
            try:
                relative_path = os.path.relpath(abs_path, self.repo_path)
            except ValueError:
                relative_path = abs_path

            return {
                "ok": True,
                "content": content,
                "path": abs_path,
                "relative_path": relative_path,
                "size": len(content),
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def _parse_liquid_includes(self, file_content: str, file_path: str, entry_base: str | None = None) -> dict[str, Any]:
        """
        Parse a Liquid file content and extract all include/render directives.

        Args:
            file_content: The content of the liquid file
            file_path: Path to the file (for resolving relative includes)
            entry_base: Optional entry-scoped resolution base

        Returns:
            {"ok": bool, "includes": list[str]}
        """
        try:
            # Patterns for include/render directives
            include_pattern = r'\{%[-]?\s*(?:include|render)\s+[\'"]([^\'"]+)[\'"]'

            includes = []

            # Find quoted includes
            for match in re.finditer(include_pattern, file_content):
                include_path = match.group(1)
                resolved = self._resolve_template_path(include_path, file_path, entry_base=entry_base)
                if resolved:
                    try:
                        relative = os.path.relpath(resolved, self.repo_path)
                    except ValueError:
                        relative = resolved
                    includes.append({
                        "raw": include_path,
                        "resolved": resolved,
                        "relative": relative,
                    })
                else:
                    includes.append({
                        "raw": include_path,
                        "resolved": None,
                        "error": "Could not resolve path"
                    })

            return {
                "ok": True,
                "includes": includes,
                "include_count": len(includes),
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}

    def build_include_tree(self, entry_point: str) -> dict[str, Any]:
        """
        Build the complete include tree starting from an entry point template.

        Args:
            entry_point: Path to the entry point template
                         (e.g., 'Service.Chat/Inference/Prompts/Templates/syndication-api/responding/prompt.liquid' or
                          'Orchestration/Extensions/WebSearch/LiquidTemplates/actions/Search.liquid')

        Returns:
            {"ok": bool, "tree": dict, "all_templates": list, "total_templates": int}
        """
        try:
            # Derive entry-scoped base for include resolution
            entry_base = self._derive_entry_base(entry_point)

            visited: set[str] = set()
            all_templates: list[dict[str, Any]] = []

            def build_tree_recursive(template_path: str, depth: int = 0) -> dict[str, Any] | None:
                # Read the file
                read_result = self._read_liquid_file(template_path, entry_base=entry_base)
                if not read_result["ok"]:
                    return {"path": template_path, "error": read_result.get("error"), "children": []}

                resolved_path = read_result["path"]

                # Avoid circular includes
                if resolved_path in visited:
                    return {"path": resolved_path, "circular": True, "children": []}

                visited.add(resolved_path)

                # Parse includes
                parse_result = self._parse_liquid_includes(read_result["content"], resolved_path, entry_base=entry_base)

                # Build template info - include content hash for caching
                content = read_result["content"]
                content_hash = f"sha256:{hashlib.sha256(content.encode('utf-8')).hexdigest()[:16]}"

                # Sanitize content preview by replacing problematic characters
                content_preview = content[:500] + "..." if len(content) > 500 else content
                content_preview = content_preview.encode('ascii', 'replace').decode('ascii')

                template_info = {
                    "path": resolved_path,
                    "relative_path": read_result.get("relative_path", template_path),
                    "name": os.path.basename(resolved_path),
                    "depth": depth,
                    "content_hash": content_hash,
                    "content_preview": content_preview,
                }
                all_templates.append(template_info)

                # Recursively process children
                children = []
                if parse_result["ok"]:
                    for include in parse_result.get("includes", []):
                        if include.get("resolved"):
                            child_tree = build_tree_recursive(include["resolved"], depth + 1)
                            if child_tree:
                                children.append(child_tree)

                return {
                    "path": resolved_path,
                    "relative_path": read_result.get("relative_path", template_path),
                    "name": os.path.basename(resolved_path),
                    "children": children,
                }

            tree = build_tree_recursive(entry_point)

            return {
                "ok": True,
                "tree": tree,
                "all_templates": all_templates,
                "total_templates": len(all_templates),
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}

    async def generate_structure_doc_from_tree(
        self,
        tree_result: dict[str, Any],
        model_client: Any,
        logger: Any = None,
        broadcast_callback: Any = None
    ) -> str:
        """
        Generate markdown structure documentation from a build_include_tree result using LLM.

        Args:
            tree_result: Result from build_include_tree containing 'tree' and 'all_templates'
            model_client: The model client to use for LLM calls
            logger: Optional logger for progress messages
            broadcast_callback: Optional callback for broadcasting progress to frontend

        Returns:
            Markdown string with complete template documentation
        """
        from autogen_core.models import UserMessage, SystemMessage

        if not tree_result.get("ok"):
            return f"# Error\n\nFailed to build tree: {tree_result.get('error')}"

        tree = tree_result.get("tree", {})
        all_templates = tree_result.get("all_templates", [])
        template_to_entry_points = tree_result.get("template_to_entry_points", {})

        # Build hierarchy string
        def build_hierarchy(node: dict, prefix: str = "", is_last: bool = True) -> list[str]:
            lines = []
            connector = "└── " if is_last else "├── "
            lines.append(f"{prefix}{connector}{node.get('name', 'unknown')}")

            children = node.get("children", [])
            child_prefix = prefix + ("    " if is_last else "│   ")
            for i, child in enumerate(children):
                is_child_last = (i == len(children) - 1)
                lines.extend(build_hierarchy(child, child_prefix, is_child_last))

            return lines

        hierarchy_lines = [tree.get("name", "prompt.liquid")]
        for i, child in enumerate(tree.get("children", [])):
            is_last = (i == len(tree.get("children", [])) - 1)
            hierarchy_lines.extend(build_hierarchy(child, "", is_last))

        hierarchy_str = "\n".join(hierarchy_lines)

        # Build template documentation sections using LLM in parallel
        total = len(all_templates)

        if logger is not None:
            logger.info(f"Inferring purpose for {total} templates in parallel (max 5 concurrent)...")
        else:
            module_logger.info("Inferring purpose for %s templates in parallel (max 5 concurrent)...", total)

        # Use semaphore to limit concurrent LLM calls to 5
        semaphore = asyncio.Semaphore(5)
        completed_count = 0
        completed_lock = asyncio.Lock()

        async def infer_with_semaphore(idx: int, template: dict[str, Any]) -> tuple[int, dict[str, Any], str]:
            """Wrapper to call LLM with semaphore and return (index, template, purpose)."""
            nonlocal completed_count
            async with semaphore:
                name = template.get("name", "unknown")
                content_preview = template.get("content_preview", "")
                purpose = await self._infer_template_purpose_llm(name, content_preview, model_client)

                # Update progress
                async with completed_lock:
                    completed_count += 1
                    # Broadcast progress every 5 templates or at milestones
                    if broadcast_callback and (completed_count % 5 == 0 or completed_count == total):
                        try:
                            await broadcast_callback(
                                "agent_message",
                                "System",
                                f"Analyzed {completed_count}/{total} templates..."
                            )
                        except Exception:
                            pass  # Don't fail on broadcast errors

                return (idx, template, purpose)

        # Create tasks for all templates
        tasks = [infer_with_semaphore(idx, template) for idx, template in enumerate(all_templates)]

        # Run all tasks concurrently with return_exceptions=True for graceful error handling
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Build doc sections from results, maintaining original order
        doc_sections = []
        for result in results:
            if isinstance(result, Exception):
                # Fallback for failed LLM calls
                doc_sections.append(f"""### unknown
- **Path:** `unknown`
- **Purpose:** Template functionality (parallel inference error: {str(result)[:50]})
""")
            else:
                idx, template, purpose = result
                name = template.get("name", "unknown")
                relative_path = template.get("relative_path", "")

                # Check if this template is shared across multiple entry points
                entry_points = template_to_entry_points.get(relative_path, [])
                shared_marker = ""
                if len(entry_points) > 1:
                    shared_marker = f" 🔗 *Shared with: {', '.join(entry_points)}*"

                section = f"""### {name}
- **Path:** `{relative_path}`{shared_marker}
- **Purpose:** {purpose}
"""
                doc_sections.append(section)

        if logger is not None:
            logger.info(f"Completed inferring purpose for {total} templates")
        else:
            module_logger.info("Completed inferring purpose for %s templates", total)

        # Combine all parts - use subsection headers for better organization
        full_doc = f"""### Template Hierarchy
```
{hierarchy_str}
```

### Template Documentation

{chr(10).join(doc_sections)}
"""
        return full_doc

    async def _infer_template_purpose_llm(
        self,
        name: str,
        content_preview: str,
        model_client: Any,
        max_retries: int = 3
    ) -> str:
        """
        Use LLM to infer template purpose from name and content preview.

        Args:
            name: Template file name
            content_preview: First 500 chars of template content
            model_client: The model client to use for LLM calls
            max_retries: Maximum number of retry attempts (default: 3)

        Returns:
            Template purpose description (one-line, <150 chars)
        """
        from autogen_core.models import UserMessage, SystemMessage
        import asyncio

        system_prompt = (
            "You are analyzing a Liquid template file to determine its purpose. "
            "Based on the template name and content preview, provide a concise one-line description "
            "of what this template does. Focus on the functional purpose, not implementation details. "
            "Keep the description under 100 characters. Do not include any prefix like 'Purpose:' - just the description."
        )

        user_prompt = f"Template name: {name}\n\nContent preview:\n{content_preview}"

        for attempt in range(max_retries):
            try:
                result = await model_client.create(
                    messages=[
                        SystemMessage(content=system_prompt),
                        UserMessage(content=user_prompt, source="user"),
                    ]
                )
                purpose = result.content.strip()
                # Ensure it's not too long
                if len(purpose) > 150:
                    purpose = purpose[:147] + "..."
                return purpose

            except Exception as e:
                error_msg = str(e)
                error_type = type(e).__name__

                if attempt < max_retries - 1:
                    # Retry for all errors with exponential backoff
                    wait_time = 2 ** attempt  # 1s, 2s, 4s
                    module_logger.warning(
                        "LLM error for %s (%s) - attempt %d/%d, retrying in %ds...",
                        name, error_type, attempt + 1, max_retries, wait_time
                    )
                    await asyncio.sleep(wait_time)
                else:
                    # Final attempt failed
                    module_logger.error(
                        "LLM failed for %s after %d attempts: %s",
                        name, max_retries, error_msg[:100]
                    )
                    return f"Template functionality (error after {max_retries} attempts)"

        # Should never reach here, but just in case
        return "Template functionality (max retries exceeded)"

    async def build_and_write_structure_doc(self, entry_point: str, model_client: Any, logger: Any = None) -> dict[str, Any]:
        """
        Build include tree and automatically generate + write structure documentation.

        This is a convenience method that combines build_include_tree,
        generate_structure_doc_from_tree, and write_structure_doc.

        Args:
            entry_point: Path to the entry point template
            model_client: The model client to use for LLM calls
            logger: Optional logger for progress messages

        Returns:
            {"ok": bool, "file_path": str, "total_templates": int}
        """
        # Check if structure doc already exists - skip generation if so
        if self.structure_doc_exists():
            msg = f"Structure doc already exists: {self._structure_doc_path}"
            if logger is not None:
                logger.info(msg)
            else:
                module_logger.info(msg)
            # Rebuild all_templates from tree for filtering cache even when doc exists
            tree_result = self.build_include_tree(entry_point)
            if tree_result.get("ok"):
                self._last_all_templates = tree_result.get("all_templates", [])
            return {
                "ok": True,
                "file_path": str(self._structure_doc_path),
                "total_templates": 0,  # Unknown, but file exists
                "size": os.path.getsize(str(self._structure_doc_path)),
                "skipped": True,
            }

        # Build tree
        tree_result = self.build_include_tree(entry_point)
        if not tree_result.get("ok"):
            return tree_result

        # Generate doc using LLM
        doc_content = await self.generate_structure_doc_from_tree(tree_result, model_client, logger)

        # Write doc
        write_result = self.write_structure_doc(doc_content)
        if not write_result.get("ok"):
            return write_result

        # Cache all_templates for condition-aware filtering
        self._last_all_templates = tree_result.get("all_templates", [])

        return {
            "ok": True,
            "file_path": write_result.get("file_path"),
            "total_templates": tree_result.get("total_templates"),
            "size": write_result.get("size"),
        }

    async def build_and_write_structure_doc_from_entries(
        self,
        liquid_entries: list[dict[str, Any]],
        model_client: Any,
        logger: Any = None,
        broadcast_callback: Any = None
    ) -> dict[str, Any]:
        """
        Build include tree from multiple liquid entry points and generate structure documentation.

        For entries with analyze_structure=True, build the full include tree (analyze all nested includes).
        For entries with analyze_structure=False, analyze only that specific file (no include tree expansion).

        Args:
            liquid_entries: List of {"path": str, "analyze_structure": bool} entries
            model_client: The model client to use for LLM calls
            logger: Optional logger for progress messages
            broadcast_callback: Optional callback for progress broadcasting (event_type, agent_name, message)

        Returns:
            {"ok": bool, "file_path": str, "total_templates": int}
        """
        # Check if structure doc already exists - skip generation if so
        if self.structure_doc_exists():
            msg = f"Structure doc already exists: {self._structure_doc_path}"
            if logger is not None:
                logger.info(msg)
            else:
                module_logger.info(msg)
            return {
                "ok": True,
                "file_path": self._structure_doc_path,
                "total_templates": 0,
                "size": os.path.getsize(self._structure_doc_path),
                "skipped": True,
            }

        all_templates = []
        analyzed_trees = []
        non_analyzed_paths = []
        seen_template_paths = set()  # Track which templates we've already documented

        # Process each entry
        for entry in liquid_entries:
            path = entry.get("path", "")
            analyze = entry.get("analyze_structure", True)

            if not path:
                continue

            if analyze:
                # Build full include tree
                tree_result = self.build_include_tree(path)
                if tree_result.get("ok"):
                    analyzed_trees.append({
                        "entry_point": path,
                        "tree": tree_result.get("tree", {}),
                        "templates": tree_result.get("all_templates", [])
                    })
                    # Track all templates from this tree
                    for template in tree_result.get("all_templates", []):
                        template_path = template.get("relative_path") or template.get("path")
                        if template_path and template_path not in seen_template_paths:
                            all_templates.append(template)
                            seen_template_paths.add(template_path)
            else:
                # Just add the path directly, no analysis
                non_analyzed_paths.append(path)

        # Generate combined documentation
        doc_sections = []

        # Add metadata header
        doc_sections.append(f"<!-- Generated: {datetime.now(timezone.utc).isoformat()} -->\n")
        doc_sections.append(f"# Liquid Template Structure Documentation\n")
        doc_sections.append(f"**Total Entry Points:** {len(analyzed_trees) + len(non_analyzed_paths)}\n")
        doc_sections.append(f"**Entry Points with Full Analysis:** {len(analyzed_trees)}\n")
        doc_sections.append(f"**Single File Analysis:** {len(non_analyzed_paths)}\n\n")
        doc_sections.append("---\n\n")

        # Add analyzed entries with full documentation
        if analyzed_trees:
            # Build a map of template paths to entry points for detecting shared templates
            template_to_entry_points = {}
            for tree_info in analyzed_trees:
                for template in tree_info["templates"]:
                    template_path = template.get("relative_path") or template.get("path")
                    if template_path:
                        if template_path not in template_to_entry_points:
                            template_to_entry_points[template_path] = []
                        template_to_entry_points[template_path].append(tree_info["entry_point"])

            total_entries = len(analyzed_trees)
            for idx, tree_info in enumerate(analyzed_trees, 1):
                entry_point = tree_info["entry_point"]

                # Count shared templates in this tree
                shared_count = 0
                for template in tree_info["templates"]:
                    template_path = template.get("relative_path") or template.get("path")
                    if template_path and len(template_to_entry_points.get(template_path, [])) > 1:
                        shared_count += 1

                doc_sections.append(f"## Entry Point {idx}: {entry_point}\n")
                doc_sections.append(f"*Full include tree analysis with {len(tree_info['templates'])} templates")
                if shared_count > 0:
                    doc_sections.append(f" ({shared_count} shared with other entry points)")
                doc_sections.append("*\n\n")

                # Generate documentation for this tree
                tree_result = {
                    "ok": True,
                    "tree": tree_info["tree"],
                    "all_templates": tree_info["templates"],
                    "template_to_entry_points": template_to_entry_points  # Pass for shared template marking
                }

                # Create a wrapper callback that adds entry point context to progress messages
                entry_broadcast_callback = None
                if broadcast_callback:
                    async def make_entry_callback(entry_idx, total):
                        async def wrapper(event_type, agent_name, message):
                            # Add entry point prefix to progress messages
                            if "Analyzed" in message and "templates" in message:
                                prefixed_message = f"Entry {entry_idx}/{total}: {message}"
                                await broadcast_callback(event_type, agent_name, prefixed_message)
                            else:
                                await broadcast_callback(event_type, agent_name, message)
                        return wrapper
                    entry_broadcast_callback = await make_entry_callback(idx, total_entries)

                doc_content = await self.generate_structure_doc_from_tree(
                    tree_result, model_client, logger, broadcast_callback=entry_broadcast_callback
                )
                doc_sections.append(doc_content)
                doc_sections.append("\n---\n\n")

        # Add non-analyzed entries (analyze single file only, no include tree)
        if non_analyzed_paths:
            doc_sections.append("## Single File Analysis\n")
            doc_sections.append(f"*{len(non_analyzed_paths)} file(s) analyzed individually without include tree traversal*\n\n")

            # Analyze each single file
            for path in non_analyzed_paths:
                # Read the file content - all paths are now repo-relative
                template_path = Path(self.repo_path) / path
                if not template_path.exists():
                    # Fallback: try templates_base for legacy short paths
                    fallback = Path(self.templates_base) / path
                    if fallback.exists():
                        template_path = fallback
                try:
                    display_path = os.path.relpath(str(template_path), self.repo_path)
                except ValueError:
                    display_path = str(template_path)
                display_path = display_path.replace("/", "\\")
                if template_path.exists():
                    try:
                        # Try reading with UTF-8 encoding
                        try:
                            content = template_path.read_text(encoding="utf-8")
                        except UnicodeDecodeError:
                            # Fallback to different encodings
                            for encoding in ["latin-1", "cp1252", "utf-16"]:
                                try:
                                    content = template_path.read_text(encoding=encoding)
                                    module_logger.warning(
                                        "File %s decoded with %s encoding (UTF-8 failed)",
                                        path, encoding
                                    )
                                    break
                                except UnicodeDecodeError:
                                    continue
                            else:
                                raise UnicodeDecodeError(
                                    "utf-8", b"", 0, 1,
                                    "Could not decode file with any common encoding"
                                )

                        # Check if file is empty
                        if not content.strip():
                            doc_sections.append(f"### {template_path.name}\n")
                            doc_sections.append(f"- **Path:** `{display_path}`\n")
                            doc_sections.append(f"- **Warning:** File is empty\n\n")
                            module_logger.warning("File %s is empty", template_path)
                            continue

                        content_preview = content[:500] if len(content) > 500 else content

                        # Use LLM to infer purpose with retry logic
                        try:
                            purpose = await self._infer_template_purpose_llm(
                                template_path.name,
                                content_preview,
                                model_client
                            )
                        except Exception as llm_error:
                            module_logger.error(
                                "LLM inference failed for %s: %s",
                                path, str(llm_error)
                            )
                            purpose = f"[Analysis failed: {str(llm_error)[:100]}]"

                        # Add to documentation
                        doc_sections.append(f"### {template_path.name}\n")
                        doc_sections.append(f"- **Path:** `{display_path}`\n")
                        doc_sections.append(f"- **Purpose:** {purpose}\n\n")

                        # Add to all_templates
                        all_templates.append({
                            "path": str(template_path),
                            "relative_path": path,
                            "name": template_path.name,
                            "depth": 0,
                        })
                    except UnicodeDecodeError as e:
                        doc_sections.append(f"### {Path(path).name}\n")
                        doc_sections.append(f"- **Path:** `{display_path}`\n")
                        doc_sections.append(f"- **Error:** Encoding error - {str(e)}\n\n")
                        module_logger.error("Encoding error reading %s: %s", template_path, e)
                    except PermissionError as e:
                        doc_sections.append(f"### {Path(path).name}\n")
                        doc_sections.append(f"- **Path:** `{display_path}`\n")
                        doc_sections.append(f"- **Error:** Permission denied\n\n")
                        module_logger.error("Permission denied reading %s: %s", template_path, e)
                    except Exception as e:
                        doc_sections.append(f"### {Path(path).name}\n")
                        doc_sections.append(f"- **Path:** `{display_path}`\n")
                        doc_sections.append(f"- **Error:** Could not read file: {str(e)}\n\n")
                        module_logger.error("Unexpected error reading %s: %s", template_path, e, exc_info=True)
                else:
                    doc_sections.append(f"### {Path(template_path).name}\n")
                    doc_sections.append(f"- **Path:** `{display_path}`\n")
                    doc_sections.append(f"- **Error:** File not found\n\n")
                    module_logger.error("File not found: %s", template_path)

        # Add shared templates summary if there are analyzed trees
        if analyzed_trees and len(analyzed_trees) > 1:
            # Build shared template map from seen_template_paths
            template_to_entry_points = {}
            for tree_info in analyzed_trees:
                for template in tree_info["templates"]:
                    template_path = template.get("relative_path") or template.get("path")
                    if template_path:
                        if template_path not in template_to_entry_points:
                            template_to_entry_points[template_path] = []
                        template_to_entry_points[template_path].append(tree_info["entry_point"])

            # Filter to only shared templates
            shared_templates = {path: entry_points for path, entry_points in template_to_entry_points.items() if len(entry_points) > 1}

            if shared_templates:
                doc_sections.append("---\n\n")
                doc_sections.append("## Shared Templates Summary\n\n")
                doc_sections.append(f"*{len(shared_templates)} template(s) are used by multiple entry points*\n\n")

                # Sort by number of entry points (most shared first)
                sorted_shared = sorted(shared_templates.items(), key=lambda x: len(x[1]), reverse=True)

                for template_path, entry_points in sorted_shared:
                    doc_sections.append(f"- **{template_path}**  \n")
                    doc_sections.append(f"  Used by {len(entry_points)} entry points: {', '.join(entry_points)}\n\n")

        # Cache all_templates for condition-aware filtering
        self._last_all_templates = all_templates

        # Combine all sections
        final_doc = "\n".join(doc_sections)

        # Write doc
        write_result = self.write_structure_doc(final_doc)
        if not write_result.get("ok"):
            return write_result

        return {
            "ok": True,
            "file_path": write_result.get("file_path"),
            "total_templates": len(all_templates),
            "size": write_result.get("size"),
        }


    def write_structure_doc(self, doc_content: str) -> dict[str, Any]:
        """
        Write the structure documentation to a markdown file.

        Args:
            doc_content: The markdown documentation content

        Returns:
            {"ok": bool, "file_path": str, "error": str (optional)}
        """
        try:
            # Ensure output directory exists
            output_dir = self._structure_doc_path.parent
            output_dir.mkdir(parents=True, exist_ok=True)

            # Add generation timestamp
            header = f"<!-- Generated: {datetime.now(timezone.utc).isoformat()} -->\n\n"
            full_content = header + doc_content

            # Write with atomic operation (write to temp file then rename)
            temp_path = self._structure_doc_path.with_suffix('.tmp')
            try:
                with open(temp_path, "w", encoding="utf-8") as f:
                    f.write(full_content)

                # Atomic rename (Windows-safe)
                if temp_path.exists():
                    if self._structure_doc_path.exists():
                        self._structure_doc_path.unlink()
                    temp_path.rename(self._structure_doc_path)

            finally:
                # Clean up temp file if it still exists
                if temp_path.exists():
                    temp_path.unlink()

            return {
                "ok": True,
                "file_path": str(self._structure_doc_path),
                "size": len(full_content),
            }
        except PermissionError as e:
            error_msg = f"Permission denied writing to {self._structure_doc_path}: {e}"
            module_logger.error(error_msg)
            return {"ok": False, "error": error_msg}
        except OSError as e:
            error_msg = f"OS error writing to {self._structure_doc_path}: {e}"
            module_logger.error(error_msg)
            return {"ok": False, "error": error_msg}
        except Exception as e:
            error_msg = f"Unexpected error writing structure doc: {e}"
            module_logger.error(error_msg, exc_info=True)
            return {"ok": False, "error": str(e)}

    async def filter_priority_templates(
        self,
        rendered_prompt: str,
        priority_templates: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """Run deterministic filtering pipeline: parse conditionals -> string match -> filter.

        Args:
            rendered_prompt: Rendered prompt text from evaluation.
            priority_templates: List of priority template dicts (with template_name, template_path).

        Returns:
            {"ok": True, "filtered_templates": [{"template_name", "template_path", "filtered_content", ...}]}
            filtered_content is None if no conditionals found, filtering was skipped, or all blocks are low confidence.
        """
        filtered_templates = []
        for pt in priority_templates:
            template_name = pt["template_name"]
            template_path = pt["template_path"]

            # Read the source file
            read_result = self._read_liquid_file(template_path)
            if not read_result.get("ok"):
                filtered_templates.append({
                    "template_name": template_name,
                    "template_path": template_path,
                    "filtered_content": None,
                    "error": f"Could not read file: {read_result.get('error')}",
                })
                continue

            source = read_result["content"]

            # Check if file has any conditionals at all
            tag_re = re.compile(r'\{%-?\s*(if|unless)\b')
            if not tag_re.search(source):
                filtered_templates.append({
                    "template_name": template_name,
                    "template_path": template_path,
                    "filtered_content": None,
                })
                continue

            # Deterministic branch matching (sequential cursor-based)
            branch_mapping = match_conditional_branches_sequential(source, rendered_prompt)

            high_confidence_blocks = [b for b in branch_mapping if b.get("confidence") == "high"]
            low_confidence_blocks = [b for b in branch_mapping if b.get("confidence") == "low"]
            module_logger.info(
                "Template %s: %d conditional blocks, %d high-confidence, %d low-confidence",
                template_name, len(branch_mapping), len(high_confidence_blocks), len(low_confidence_blocks),
            )
            for b in low_confidence_blocks:
                module_logger.info(
                    "  Low-confidence block lines %s-%s: method=%s, active_branch=%s",
                    b.get("line_start", "?"), b.get("line_end", "?"),
                    b.get("match_method", "?"), b.get("active_branch", "?"),
                )

            # If all blocks are low confidence, skip filtering (preserve all branches)
            if not high_confidence_blocks:
                module_logger.info(
                    "All conditional blocks in %s are low confidence, skipping filter",
                    template_name,
                )
                filtered_templates.append({
                    "template_name": template_name,
                    "template_path": template_path,
                    "filtered_content": None,
                })
                continue

            # Build filtered content
            filter_result = build_filtered_content(source, branch_mapping)
            filtered_content = filter_result["filtered_content"]

            # If filtering removed ALL content (e.g. entire file wrapped in an
            # inactive conditional), fall back to None so the sidecar is not
            # written and the optimizer reads the original source instead.
            if not filtered_content or not filtered_content.strip():
                module_logger.info(
                    "Template %s: filtered content is empty (all blocks inactive), skipping sidecar",
                    template_name,
                )
                filtered_templates.append({
                    "template_name": template_name,
                    "template_path": template_path,
                    "filtered_content": None,
                })
                continue

            filtered_templates.append({
                "template_name": template_name,
                "template_path": template_path,
                "filtered_content": filtered_content,
                "active_line_ranges": filter_result["active_line_ranges"],
                "line_map": filter_result["line_map"],
            })

        return {"ok": True, "filtered_templates": filtered_templates}

    def embed_filtered_content_in_priority_templates(
        self,
        priority_templates_path: str,
        filtered_templates: list[dict[str, Any]],
    ) -> None:
        """Append filtered content sections to existing priority_templates.md.

        Args:
            priority_templates_path: Path to the priority_templates.md file.
            filtered_templates: Output from filter_priority_templates().
        """
        filtered_by_name: dict[str, dict] = {}
        for ft in filtered_templates:
            filtered_by_name[ft["template_name"]] = {
                "content": ft.get("filtered_content"),
                "ranges": ft.get("active_line_ranges", []),
            }

        with open(priority_templates_path, "r", encoding="utf-8") as f:
            content = f.read()

        for template_name, data in filtered_by_name.items():
            if data["content"] is None:
                continue

            # Find the section for this template: ## N. template_name
            # Match the header line and all subsequent "- **...**" lines
            pattern = re.compile(
                rf"(## \d+\. {re.escape(template_name)}\n(?:- \*\*.*\n)*)",
                re.MULTILINE,
            )
            match = pattern.search(content)
            if match:
                insertion = (
                    f"- **Active content (for modification):**\n"
                    f"\n```liquid\n{data['content']}\n```\n"
                )
                insert_pos = match.end()
                content = content[:insert_pos] + insertion + content[insert_pos:]

        with open(priority_templates_path, "w", encoding="utf-8") as f:
            f.write(content)

    def write_filtered_sidecars(
        self,
        filtered_templates: list[dict[str, Any]],
    ) -> None:
        """Write sidecar files and manifest for transparent optimizer read/write.

        Creates .filtered/<relative_path> sidecar files containing the filtered
        content, and .filtered/_manifest.json with line_map and active_line_ranges
        for each filtered template.

        Args:
            filtered_templates: Output from filter_priority_templates().
        """
        filtered_dir = Path(self.repo_path) / ".filtered"

        manifest: dict[str, Any] = {}

        for ft in filtered_templates:
            fc = ft.get("filtered_content")
            if fc is None:
                continue

            template_path = ft["template_path"]
            line_map = ft.get("line_map", [])
            active_ranges = ft.get("active_line_ranges", [])

            # Write sidecar file
            sidecar_path = filtered_dir / template_path
            sidecar_path.parent.mkdir(parents=True, exist_ok=True)
            sidecar_path.write_text(fc, encoding="utf-8")

            # Add to manifest
            manifest[template_path] = {
                "active_line_ranges": active_ranges,
                "line_map": line_map,
            }

        if manifest:
            filtered_dir.mkdir(parents=True, exist_ok=True)
            manifest_path = filtered_dir / "_manifest.json"
            manifest_path.write_text(
                json.dumps(manifest, indent=2),
                encoding="utf-8",
            )

    async def analyze_bad_cases_for_templates_llm(
        self,
        bad_case_analysis: str | list[dict[str, Any]],
        model_client: Any,
        structure_doc: str | None = None,
        output_dir: str | None = None,
    ) -> dict[str, Any]:
        """
        Analyze bad case clusters using LLM to identify 2-3 priority templates for optimization.

        This function uses an LLM to analyze failure patterns and map them to specific
        liquid templates that should be modified.

        Args:
            bad_case_analysis: JSONL string or list of cluster dictionaries from BadCaseAnalyzer
            model_client: The model client to use for LLM calls
            structure_doc: Optional structure doc content (will read from file if not provided)
            output_dir: Optional directory to save priority_templates.md (defaults to store_dir)

        Returns:
            {
                "ok": bool,
                "priority_templates_path": str,
                "priority_templates": [...],
                "analysis_summary": str,
                "total_clusters_analyzed": int
            }
        """
        from autogen_core.models import UserMessage

        try:
            # Parse bad case analysis - handle multiple formats
            content = bad_case_analysis.strip()

            if not content:
                return {"ok": False, "error": f"No valid content found in bad case analysis. Content preview: {str(bad_case_analysis)[:200]}"}

            # Read structure doc if not provided
            if structure_doc is None:
                structure_doc = self.read_structure_doc()

            if not structure_doc:
                return {"ok": False, "error": "Structure documentation not available"}

            # Build LLM prompt
            llm_prompt = f"""You are an expert in Liquid template analysis for prompt optimization.

## Task
Analyze the following bad case clusters from evaluation results and identify which 2-3 Liquid templates should be prioritized for optimization.

## Bad Case Clusters (from BadCaseAnalyzer)
```json
{content}
```

## Available Templates (from Structure Documentation)
{structure_doc}

## Instructions
1. Read each bad case cluster carefully - pay attention to the "pattern", "root_cause", and "suggestion" fields
2. Map each cluster to the most relevant template(s) based on:
   - The type of issue (formatting, citation, persona, safety, etc.)
   - The "Purpose" and "DO modify here" guidance in the structure doc
3. Select 2-3 templates that would have the HIGHEST IMPACT on fixing the identified issues
4. For each selected template, explain WHY it should be modified and WHAT specific changes to consider

## Output Format
Respond with a JSON object (no markdown code block, just raw JSON):
{{
  "priority_templates": [
    {{
      "template_name": "example.liquid",
      "reason": "Why this template is prioritized based on bad case analysis",
      "related_clusters": ["cluster_id_1", "cluster_id_2"],
      "suggested_changes": ["Specific change 1", "Specific change 2"]
    }}
  ],
  "analysis_summary": "Brief summary of the analysis"
}}

Important:
- Only select templates that are mentioned in the Structure Documentation
- Focus on templates where changes would address multiple bad case clusters
- Be specific about what changes should be made
- IMPORTANT: Entry point paths (shown as "## Entry Point N: ...") are NOT templates you can select. Only select templates from the "### template_name" sections under Template Documentation. Use the exact template filename (e.g., "idiolect.liquid") as template_name, not the entry point path."""

            # Call LLM
            messages = [UserMessage(content=llm_prompt, source="liquid_structure_agent")]
            response = await model_client.create(messages)

            # Parse LLM response
            response_text = response.content
            if isinstance(response_text, list):
                response_text = response_text[0].text if response_text else ""

            # Try to extract JSON from response
            try:
                # Handle potential markdown code blocks
                if "```json" in response_text:
                    json_start = response_text.find("```json") + 7
                    json_end = response_text.find("```", json_start)
                    response_text = response_text[json_start:json_end].strip()
                elif "```" in response_text:
                    json_start = response_text.find("```") + 3
                    json_end = response_text.find("```", json_start)
                    response_text = response_text[json_start:json_end].strip()

                llm_result = json.loads(response_text)
            except json.JSONDecodeError as e:
                return {"ok": False, "error": f"Failed to parse LLM response as JSON: {e}. Response: {response_text[:500]}"}

            # Extract template info from structure doc for each priority template
            def find_template_info(template_name: str) -> dict[str, str]:
                """Extract Path, Purpose, DO/DON'T modify from structure doc."""
                template_base = template_name.replace('.liquid', '')

                # Try to find the section for this template
                # Pattern matches "### template_name" until next "###" or end
                section_pattern = rf"###\s+{re.escape(template_name)}.*?(?=\n###|\Z)"
                section_match = re.search(section_pattern, structure_doc, re.DOTALL | re.IGNORECASE)

                info = {"path": "", "purpose": "", "do_modify": "", "dont_modify": ""}

                if section_match:
                    section = section_match.group(0)

                    # Extract Path - handle both "**Path:**" and "Path:" formats
                    path_match = re.search(r"(?:\*\*)?Path:(?:\*\*)?\s*`?([^\n`]+)`?", section)
                    if path_match:
                        info["path"] = path_match.group(1).strip()

                    # Extract Purpose - handle both bold and non-bold formats
                    purpose_match = re.search(r"(?:\*\*)?Purpose:(?:\*\*)?\s*([^\n]+)", section)
                    if purpose_match:
                        info["purpose"] = purpose_match.group(1).strip()

                    # Extract DO modify - try multiple patterns
                    do_match = re.search(r"(?:\*\*)?DO modify(?: here)?:(?:\*\*)?\s*([^\n]+)", section, re.IGNORECASE)
                    if do_match:
                        info["do_modify"] = do_match.group(1).strip()
                    else:
                        # Try "Modification guidance" as alternative
                        mod_match = re.search(r"(?:\*\*)?Modification guidance:(?:\*\*)?\s*([^\n]+)", section, re.IGNORECASE)
                        if mod_match:
                            info["do_modify"] = mod_match.group(1).strip()

                    # Extract DON'T modify
                    dont_match = re.search(r"(?:\*\*)?DON'T modify(?: here)?:(?:\*\*)?\s*([^\n]+)", section, re.IGNORECASE)
                    if dont_match:
                        info["dont_modify"] = dont_match.group(1).strip()

                # Fallback: search entire doc for path if not found in section
                if not info["path"]:
                    # Look for path line containing the template name - handle both bold and non-bold formats
                    path_pattern = rf"(?:\*\*)?Path:(?:\*\*)?\s*`?([^\n`]*{re.escape(template_base)}[^\n`]*)`?"
                    path_match = re.search(path_pattern, structure_doc, re.IGNORECASE)
                    if path_match:
                        info["path"] = path_match.group(1).strip()
                    else:
                        info["path"] = f"(path not found for {template_name})"

                return info

            # Build set of known template names from ### headers in structure doc
            known_templates = set(re.findall(r"^###\s+(\S+\.liquid)", structure_doc, re.MULTILINE))

            # Enrich priority templates with structure doc info
            priority_templates = []
            for template in llm_result.get("priority_templates", []):
                template_name = template.get("template_name", "")

                # Validate template_name against known ### sections
                if template_name not in known_templates:
                    # LLM may have returned an entry point path (e.g., "actions/search_finance.liquid")
                    # Try matching by basename
                    basename = os.path.basename(template_name)
                    if basename in known_templates:
                        module_logger.warning(
                            "Normalized template_name from '%s' to '%s' (matched by basename)",
                            template_name, basename
                        )
                        template_name = basename
                    else:
                        module_logger.warning(
                            "Skipping unknown template '%s' - not found in structure doc ### sections",
                            template_name
                        )
                        continue

                template_info = find_template_info(template_name)
                priority_templates.append({
                    "template_name": template_name,
                    "template_path": template_info["path"],
                    "purpose": template_info["purpose"],
                    "do_modify": template_info["do_modify"],
                    "dont_modify": template_info["dont_modify"],
                    "reason": template.get("reason", ""),
                    "clusters": template.get("related_clusters", []),
                    "suggested_changes": template.get("suggested_changes", []),
                })

            # Generate markdown file for optimizer
            md_lines = [
                f"<!-- Generated: {datetime.now(timezone.utc).isoformat()} -->",
                "",
                "# Priority Templates for Optimization",
                "",
                f" There are {len(priority_templates)} failure clusters from LLM analysis, the following templates should be prioritized for optimization:",
                "",
            ]

            for i, template in enumerate(priority_templates, 1):
                md_lines.append(f"## {i}. {template['template_name']}")
                md_lines.append(f"- **Path:** {template['template_path']}")
                if template['purpose']:
                    md_lines.append(f"- **Purpose:** {template['purpose']}")
                if template['do_modify']:
                    md_lines.append(f"- **DO modify here:** {template['do_modify']}")
                if template['dont_modify']:
                    md_lines.append(f"- **DON'T modify here:** {template['dont_modify']}")
                md_lines.append(f"- **Why prioritized:** {template['reason']}")
                if template['clusters']:
                    md_lines.append(f"- **Related clusters:** {', '.join(template['clusters'])}")
                if template['suggested_changes']:
                    md_lines.append("- **Suggested changes:**")
                    for suggestion in template['suggested_changes']:
                        md_lines.append(f"  - {suggestion}")
                md_lines.append("")

            # Add summary section
            md_lines.append("## Summary")
            summary = llm_result.get("analysis_summary", "")
            if not summary:
                summary = f"Analyzed {len(content)} failure clusters. "
                summary += f"Identified {len(priority_templates)} priority templates for optimization: "
                summary += ", ".join([t["template_name"] for t in priority_templates])
            md_lines.append(summary)

            # Write to file - use output_dir if provided, otherwise store_dir
            save_dir = str(output_dir) if output_dir else self.store_dir
            os.makedirs(save_dir, exist_ok=True)
            priority_templates_path = os.path.join(save_dir, "priority_templates.md")
            with open(priority_templates_path, "w", encoding="utf-8") as f:
                f.write("\n".join(md_lines))

            return {
                "ok": True,
                "priority_templates_path": str(priority_templates_path),
                "priority_templates": priority_templates,
                "analysis_summary": summary,
                "total_clusters_analyzed": len(priority_templates),
            }

        except Exception as e:
            return {"ok": False, "error": str(e)}
