"""
Transform plain-string locator files into the agent-friendly format.

Input (migrated raw format):
  export const dashboardLocators = {
    userAvatarButton: "navbar_avatar-btn",
    saveButton:        'button.save-button',
  };

Output (agent-friendly):
  export const dashboardLocators = {
    userAvatarButton: { selector: "navbar_avatar-btn",   intent: "user avatar button",   stability: 80 },
    saveButton:        { selector: "button.save-button", intent: "save button",           stability: 30 },
  } as const;
  export type DashboardLocatorKey = keyof typeof dashboardLocators;

Arrow-function / method-shorthand entries are kept unchanged because they
produce dynamic selectors that cannot be statically registered.

Text-only values (display strings, URLs, error messages) are also kept as-is.
"""
from __future__ import annotations
import re

from stlc_agents.shared import heal_intel as _heal_intel


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _camel_to_words(name: str) -> str:
    """userAvatarButton → 'user avatar button'.

    Words are normalised through the heal_intel intent-keyword allowlist —
    out-of-vocabulary tokens are still preserved (a migrated `userAvatarButton`
    still becomes "user avatar button" so the human reading the registered
    locator can recognise it), but the resulting string is what
    `LocatorHealer._intentShapeFilter` consumes, so the migration step is
    using the same vocabulary the runtime healer is trained on.
    """
    s = re.sub(r"([A-Z])", r" \1", name)
    return s.strip().lower()


def _has_descendant_combinator(selector: str) -> bool:
    """True if `selector` joins two simple selectors with a combinator.

    A descendant combinator is just a space (`#search-pagination button`),
    but spaces also appear *inside* attribute selectors (`[aria-label="Save
    all"]`), quoted text, and `:has(...)` / `:not(...)` argument lists where
    they are NOT structural. Those are not combinators, so we ignore any
    space that sits inside `[]`, `()`, or quotes. The explicit combinators
    `> + ~` always count.
    """
    depth = 0          # nesting of [] and ()
    quote = ""         # current quote char, "" when outside a string
    for ch in selector:
        if quote:
            if ch == quote:
                quote = ""
            continue
        if ch in "\"'":
            quote = ch
        elif ch in "[(":
            depth += 1
        elif ch in "])":
            depth = max(0, depth - 1)
        elif depth == 0 and (ch in ">+~" or ch == " "):
            return True
    return False


def _infer_stability(selector: str) -> int:
    """0-100 confidence score for a selector. Drives healer chain ordering.

    The base scoring (data-testid → 100, id → 80, etc.) mirrors what
    `agent_playwright_generator._score_selector` returns so a migrated locator
    has the same stability tier the live healer would assign. Anything not
    covered there falls back to the legacy fixed bumps for
    `[dusk=...]` / bare test-id strings / class+combinator selectors.

    A compound selector that joins simple selectors with a combinator
    (descendant space, `>`, `+`, `~`) is brittle regardless of how stable its
    leading part looks — `#id button` breaks the moment the DOM is restructured
    — so it is always capped at the combinator tier (30), never the ID/testid
    tier its prefix would otherwise earn.
    """
    if _has_descendant_combinator(selector):
        return 30
    # Walk the heal_intel `*` order — same logic as _score_selector but
    # shape-aware for migration: we don't know the tag yet so we use the
    # generic order which closely matches the historical hardcoded values.
    order = _heal_intel.attribute_order_for("*")
    for idx, attr in enumerate(order):
        if attr == "id":
            if selector.startswith("#") or "[id=" in selector:
                return max(30, 100 - idx * 10)
        elif attr == "role+text":
            if '[role="' in selector or selector.startswith("role="):
                return max(30, 100 - idx * 10)
        elif attr == "href-shape":
            if "[href" in selector:
                return max(30, 100 - idx * 10)
        else:
            if f"[{attr}" in selector or f"{attr}=" in selector:
                return max(30, 100 - idx * 10)
    # Bare test-id strings (no CSS operators) — same fast-path as before.
    if re.match(r"^[a-z_][\w-]*$", selector):
        return 80
    # Legacy non-allowlisted but project-stable attributes.
    if "[dusk=" in selector or "[data-" in selector:
        return 70
    if selector.startswith(".") or any(c in selector for c in ">+~"):
        return 30
    return 50


def _is_pure_selector(value: str) -> bool:
    """
    Return True only for values that look like CSS / test-id selectors.
    Excludes: URLs, plain text labels, error messages.
    """
    if value.startswith("http://") or value.startswith("https://"):
        return False
    # Contains a space AND no CSS-selector chars → probably display text
    if " " in value and not any(c in value for c in "#.[>+~:=*^$|@(\"'"):
        return False
    return True


# A bare test-id NAME (no CSS metacharacters) — the value stored when the
# source consumed the locator via getByTestId(name) / waitAndClickByTestId(name)
# rather than page.locator(cssSelector). At runtime the healer feeds the stored
# selector to page.locator(...), which has no test-id engine — so a bare name
# resolves as a CSS *element* selector (`<navbar_avatar-btn>`) and matches
# nothing. Wrapping it as `[data-testid="name"]` reproduces what getByTestId
# did in the source.
_BARE_TESTID_NAME_RE = re.compile(r"^[A-Za-z][\w-]*$")


def _maybe_wrap_testid(value: str, key: str, testid_keys: set[str] | None) -> str:
    """Wrap a bare test-id name as `[data-testid="..."]` when the source
    consumed this key via a *ByTestId call. Leaves already-formed selectors
    (anything with CSS metacharacters) and non-testid keys untouched."""
    if not testid_keys or key not in testid_keys:
        return value
    if not _BARE_TESTID_NAME_RE.match(value):
        # Already a selector (`button[value='x']`), or otherwise not a bare
        # name — don't double-wrap; the source value already resolves.
        return value
    return f'[data-testid="{value}"]'


# `key: "value"` inside an xxxLocators block (single-line string entry). Used to
# decide which keys _maybe_wrap_testid actually rewrote, so callers can keep
# call-site consumption consistent with the stored selector form.
_LOC_BLOCK_OPEN = re.compile(r"export\s+const\s+(\w+Locators)\s*=\s*\{")
_LOC_STRING_ENTRY = re.compile(r"^\s*(\w+)\s*:\s*(['\"])(.*?)\2\s*,?\s*$", re.MULTILINE)


def compute_wrapped_keys(
    source_content: str,
    testid_keys_by_object: dict[str, set[str]] | None,
) -> dict[str, set[str]]:
    """Return the `{object: {keys}}` that `transform_locator_file` will rewrite
    to the `[data-testid="..."]` form for this source file.

    A key is wrapped iff it is test-id-consumed AND its source value is a bare
    name (so `_maybe_wrap_testid` changes it). Once wrapped, the stored
    `.selector` is a CSS string — so surviving `getByTestId(<obj>.<key>.selector)`
    call sites must be rewritten to `.locator(...)` to stay consistent (a
    `[data-testid="x"]` string is invalid as a getByTestId argument). This lets
    the injector know exactly which call sites to convert.
    """
    if not testid_keys_by_object:
        return {}
    wrapped: dict[str, set[str]] = {}
    for bm in _LOC_BLOCK_OPEN.finditer(source_content):
        obj = bm.group(1)
        testid_keys = testid_keys_by_object.get(obj)
        if not testid_keys:
            continue
        # Bound the scan to THIS block's body via brace-depth counting so keys
        # from a later xxxLocators block aren't attributed to this object.
        depth = 1
        j = bm.end()
        n = len(source_content)
        while j < n and depth > 0:
            c = source_content[j]
            if c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
            j += 1
        body = source_content[bm.end():j - 1]
        for em in _LOC_STRING_ENTRY.finditer(body):
            key, value = em.group(1), em.group(3)
            if key in testid_keys and _BARE_TESTID_NAME_RE.match(value):
                wrapped.setdefault(obj, set()).add(key)
    return wrapped


def _quote(value: str, original_quote: str) -> str:
    """Pick inner quote avoiding the same char that appears in value."""
    if '"' in value and "'" not in value:
        return "'"
    if "'" in value and '"' not in value:
        return '"'
    # Prefer double-quotes; escape would be needed only if value has both — rare.
    return '"'


# ---------------------------------------------------------------------------
# Per-object body transformer
# ---------------------------------------------------------------------------

def _transform_body(
    block_name: str,
    body: str,
    testid_keys: set[str] | None = None,
) -> tuple[str, list[str]]:
    """
    Transform the body text (content between the outer { }) of a locator object.
    Returns (transformed_body, list_of_change_descriptions).

    `testid_keys` is the set of keys (in this block) that the source consumed
    via a *ByTestId call; their bare-name values are rewritten to the
    `[data-testid="..."]` form so the migrated selector resolves the way the
    original getByTestId(...) did.
    """
    changes: list[str] = []
    lines = body.split("\n")
    result: list[str] = []
    i = 0

    while i < len(lines):
        raw = lines[i]
        stripped = raw.rstrip()

        # ── Pattern 1: single-line string property ──────────────────────────
        # `  key: "value",` or `  key: 'value',`
        m1 = re.match(r"^( {2,})(\w+)\s*:\s*(['\"])(.*?)\3,?\s*$", stripped)
        if m1:
            indent, key, q, value = m1.group(1), m1.group(2), m1.group(3), m1.group(4)
            if _is_pure_selector(value):
                value     = _maybe_wrap_testid(value, key, testid_keys)
                intent    = _camel_to_words(key)
                stability = _infer_stability(value)
                iq        = _quote(value, q)
                result.append(
                    f"{indent}{key}: {{ selector: {iq}{value}{iq},"
                    f' intent: "{intent}", stability: {stability} }},'
                )
                changes.append(f"Registered locator: {block_name}.{key}")
                i += 1
                continue

        # ── Pattern 2: multi-line string property ────────────────────────────
        # `  key:\n    "value",`
        m2 = re.match(r"^( {2,})(\w+)\s*:\s*$", stripped)
        if m2 and (i + 1) < len(lines):
            next_line = lines[i + 1].strip()
            mval = re.match(r"^(['\"])(.*?)\1,?\s*$", next_line)
            if mval:
                indent, key = m2.group(1), m2.group(2)
                q, value    = mval.group(1), mval.group(2)
                if _is_pure_selector(value):
                    value     = _maybe_wrap_testid(value, key, testid_keys)
                    intent    = _camel_to_words(key)
                    stability = _infer_stability(value)
                    iq        = _quote(value, q)
                    result.append(
                        f"{indent}{key}: {{ selector: {iq}{value}{iq},"
                        f' intent: "{intent}", stability: {stability} }},'
                    )
                    changes.append(f"Registered multi-line locator: {block_name}.{key}")
                    i += 2
                    continue

            # ── Pattern 3: multi-line string-concat property ─────────────────
            # ```
            #   key:
            #     "a" +
            #     "b" +
            #     "c",
            # ```
            # Collect adjacent quoted strings joined by `+` until we hit the
            # entry's trailing comma. Emits one `{ selector: "a" + "b" + "c",
            # intent, stability }` entry. Stability is computed from the
            # concatenated string so multi-line selectors get the same score
            # they'd receive if written on one line.
            indent, key = m2.group(1), m2.group(2)
            concat_re = re.compile(r"^(['\"])(.*?)\1\s*(\+)?\s*$")
            parts: list[tuple[str, str]] = []  # (quote, value)
            j = i + 1
            done = False
            while j < len(lines):
                seg = lines[j].strip().rstrip(",")
                terminates_here = lines[j].rstrip().endswith(",")
                seg_m = concat_re.match(seg)
                if not seg_m:
                    break
                parts.append((seg_m.group(1), seg_m.group(2)))
                j += 1
                if terminates_here:
                    done = True
                    break
            if done and len(parts) >= 2:
                joined_value = "".join(p[1] for p in parts)
                if _is_pure_selector(joined_value):
                    intent    = _camel_to_words(key)
                    stability = _infer_stability(joined_value)
                    # Preserve the original quoted-string form so authors keep
                    # their choice of quote and line-by-line layout signal —
                    # the runtime concatenation is identical.
                    rendered  = " + ".join(f"{q}{v}{q}" for q, v in parts)
                    result.append(
                        f"{indent}{key}: {{ selector: {rendered},"
                        f' intent: "{intent}", stability: {stability} }},'
                    )
                    changes.append(f"Registered concat locator: {block_name}.{key}")
                    i = j
                    continue

        result.append(raw)
        i += 1

    return "\n".join(result), changes


# ---------------------------------------------------------------------------
# File-level transformer
# ---------------------------------------------------------------------------

# Matches the opening of a locator constant:
#   export const xxxLocators = {
_BLOCK_OPEN_RE = re.compile(
    r"(export\s+const\s+(\w+Locators)\s*=\s*\{)[ \t]*\n",
    re.MULTILINE,
)

# After transformation we append:  } as const;\nexport type ...
_AS_CONST_RE    = re.compile(r"\}\s*;?\s*$")
_TYPE_EXPORT_RE = re.compile(r"export type \w+LocatorKey = keyof typeof \w+Locators;")


# A *ByTestId call whose argument is the locator reference DIRECTLY or as a
# pure template wrapper — both equivalent to passing the bare id:
#   getByTestId(obj.key)            waitAndClickByTestId(obj.key)
#   getByTestId(`${obj.key}`)       (pure wrapper, no suffix)
# Wrapping the stored value to `[data-testid="..."]` is correct for these.
_TESTID_REF_DIRECT_RE = re.compile(
    r"""\w*ByTestId\(\s*(\w+Locators)\.(\w+)\s*\)""")
_TESTID_REF_PUREWRAP_RE = re.compile(
    r"""\w*ByTestId\(\s*`\$\{\s*(\w+Locators)\.(\w+)\s*\}`\s*\)""")

# A template literal that COMPOSES the locator into a larger test-id by
# concatenating a test-id continuation immediately after `${obj.key}`:
#   `${obj.key}_${i}`     `${obj.key}_span_0`     `${prefix}${moduleName}`
# The continuation must look like a test-id fragment — `_`, a letter/digit, or
# another `${...}` interpolation — NOT prose. This deliberately EXCLUDES error
# messages like `${obj.key} is not visible` (a space after `}` means it's a
# string, not a test-id base). Composed bases must stay bare names, so wrapping
# them to `[data-testid="..."]` (which would break the concatenation) is
# suppressed for these keys.
_TESTID_BASE_COMPOSED_RE = re.compile(
    r"""`\$\{\s*(\w+Locators)\.(\w+)\s*\}(?:[A-Za-z0-9_-]|\$\{)""")


def collect_testid_consumed_keys(source_text: str) -> dict[str, set[str]]:
    """Scan a source .ts file for locators consumed DIRECTLY via a *ByTestId
    call (bare or pure-`${...}` wrapper).

    Returns {objectName: {key, ...}} — keys whose stored value should be wrapped
    as `[data-testid="..."]`. Keys used only as a COMPOSED test-id base
    (`getByTestId(`${obj.key}_${i}`)`) are NOT returned here and additionally
    EXCLUDED (see _composed_testid_bases), because their value must stay a bare
    name for the concatenation to produce a valid test-id. Cross-file: callers
    union over the whole source tree before transforming locator files.
    """
    out: dict[str, set[str]] = {}
    for rx in (_TESTID_REF_DIRECT_RE, _TESTID_REF_PUREWRAP_RE):
        for m in rx.finditer(source_text):
            out.setdefault(m.group(1), set()).add(m.group(2))
    # Subtract any key used as a composed base anywhere in this file — a single
    # composed usage makes wrapping unsafe regardless of other direct uses.
    composed = _composed_testid_bases(source_text)
    for obj, keys in composed.items():
        if obj in out:
            out[obj] -= keys
            if not out[obj]:
                del out[obj]
    return out


def _composed_testid_bases(source_text: str) -> dict[str, set[str]]:
    """{object: {keys}} used as a COMPOSED test-id base (`${obj.key}<suffix>`)."""
    out: dict[str, set[str]] = {}
    for m in _TESTID_BASE_COMPOSED_RE.finditer(source_text):
        out.setdefault(m.group(1), set()).add(m.group(2))
    return out


def transform_locator_file(
    content: str,
    testid_keys_by_object: dict[str, set[str]] | None = None,
) -> tuple[str, list[str]]:
    """
    Transform all locator-constant blocks in a file.
    Returns (transformed_content, changes).

    `testid_keys_by_object` maps each xxxLocators object name to the set of its
    keys that the source consumed via a *ByTestId call. Bare-name values for
    those keys are rewritten to `[data-testid="..."]` (see _maybe_wrap_testid).
    """
    if "selector:" in content and "intent:" in content:
        # Already in agent-friendly format — idempotent
        return content, []

    all_changes: list[str] = []
    output_parts: list[str] = []
    cursor = 0

    for m in _BLOCK_OPEN_RE.finditer(content):
        block_start = m.start()
        body_start  = m.end()        # first char after the opening `{\n`
        block_name  = m.group(2)

        # Find the matching closing brace with depth counting
        depth = 1
        j     = body_start
        while j < len(content) and depth > 0:
            c = content[j]
            if c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
            j += 1
        # j now points to one past the closing `}`
        body   = content[body_start : j - 1]   # exclude the final `}`
        suffix = content[j - 1 :]               # `}...` to end (or next block)

        # Transform the body
        block_testid_keys = (
            (testid_keys_by_object or {}).get(block_name)
        )
        new_body, changes = _transform_body(block_name, body, block_testid_keys)
        all_changes.extend(changes)

        # Build a pascal-case type name: dashboardLocators → Dashboard
        pascal = block_name[0].upper() + block_name[1:]
        pascal = pascal.replace("Locators", "")
        type_line = f"export type {pascal}LocatorKey = keyof typeof {block_name};"

        # Everything before this block
        output_parts.append(content[cursor:block_start])

        # Reconstructed block
        output_parts.append(m.group(1) + "\n")    # `export const xxx = {`
        output_parts.append(new_body)
        # Close: `} as const;`
        output_parts.append("} as const;\n")
        if changes:  # only add type export if we actually transformed something
            output_parts.append(type_line + "\n")

        # Skip over the original closing `};` in suffix
        # suffix starts with `}` — find the end of `};` or `}`
        tail_m = re.match(r"\}\s*;?\s*\n?", suffix)
        tail_skip = tail_m.end() if tail_m else 1
        cursor = j - 1 + tail_skip

    # Append remaining file content (after last block)
    output_parts.append(content[cursor:])

    return "".join(output_parts), all_changes


# ---------------------------------------------------------------------------
# Monolith split — one source file containing multiple xxxPageLocators consts
# becomes one file per const + a barrel. Stops a single source touchpoint from
# turning every selector edit into a merge conflict.
# ---------------------------------------------------------------------------

def _const_to_kebab(const_name: str) -> str:
    """`memberEditPageLocators` → `member-edit`. Mirrors mapper._to_kebab's
    suffix stripping so the produced filenames stay consistent with the rest
    of the Helix layout."""
    stem = const_name
    for suffix in ("Locators", "Page", "PageObject"):
        while stem.endswith(suffix):
            stem = stem[: -len(suffix)]
    stem = re.sub(r"([A-Z])", r"-\1", stem).lstrip("-").lower()
    stem = re.sub(r"[_\s]+", "-", stem)
    stem = re.sub(r"-{2,}", "-", stem).strip("-")
    return stem or "unknown"


def split_locator_file_into_pages(
    content: str,
) -> tuple[list[tuple[str, str]], str] | None:
    """
    If `content` declares ≥2 top-level `export const xxxPageLocators` blocks,
    return (per_page_files, barrel) where:
      - per_page_files is [(filename_stem, file_content), ...]
        e.g. [("login", "<const + type>\\n"), ("member-edit", ...)]
        filename_stem is the kebabed const name; callers append `.locators.ts`.
      - barrel is the `index.ts` content re-exporting every per_page file.

    Returns None when there are 0 or 1 blocks (no monolith to split).

    Pre-condition: `content` is already in agent-friendly format (i.e. has
    been through `transform_locator_file` — every block is closed with
    `} as const;` and followed by its `export type … = keyof typeof …`).

    Imports, top-of-file doc comments and arbitrary helper code BETWEEN
    blocks are intentionally NOT copied into the per-page files. The migrator
    only splits when every block is a self-contained const + type pair, which
    is the shape `transform_locator_file` always produces. Anything else
    falls back to the single-file path.
    """
    blocks = list(_BLOCK_OPEN_RE.finditer(content))
    if len(blocks) < 2:
        return None

    per_page: list[tuple[str, str]] = []
    for i, m in enumerate(blocks):
        const_name = m.group(2)
        # Pull the leading doc comment if there's a /** … */ directly above
        # (no blank line between). This keeps per-page provenance attached
        # to each const when it lands in its own file.
        start = m.start()
        leading_comment_start = start
        head_slice = content[: start]
        head_trimmed = head_slice.rstrip()
        if head_trimmed.endswith("*/"):
            comment_open = head_trimmed.rfind("/**")
            if comment_open >= 0:
                between = head_slice[comment_open:].rstrip()
                # Only attach if there's no blank line between */ and export
                if "\n\n" not in head_slice[comment_open : start]:
                    leading_comment_start = comment_open

        # End of this block = body close + the immediately-following type
        # export line (if present). Easiest way: take everything up to the
        # next block's leading comment / export, or EOF for the last block.
        if i + 1 < len(blocks):
            next_start = blocks[i + 1].start()
            # Walk back past any blank lines and a /** comment that belongs
            # to the NEXT block so it doesn't get glued onto this one.
            tail_slice = content[: next_start]
            tail_trimmed = tail_slice.rstrip()
            if tail_trimmed.endswith("*/"):
                comment_open = tail_trimmed.rfind("/**")
                if comment_open >= 0 and "\n\n" not in tail_slice[comment_open : next_start]:
                    next_start = comment_open
            end = next_start
        else:
            end = len(content)

        block_text = content[leading_comment_start : end].rstrip() + "\n"
        per_page.append((_const_to_kebab(const_name), block_text))

    barrel_lines = ["// Auto-generated barrel — re-exports every per-page locator module."]
    for stem, _ in sorted(per_page, key=lambda p: p[0]):
        barrel_lines.append(f'export * from "./{stem}.locators";')
    barrel = "\n".join(barrel_lines) + "\n"

    return per_page, barrel


# ---------------------------------------------------------------------------
# Post-migration: fix selector accesses in page objects / step files
# ---------------------------------------------------------------------------

def collect_object_keys(locator_file_content: str) -> dict[str, set[str]]:
    """
    Parse an already-transformed locator file and return a mapping of
    { locator_name → set_of_keys_that_are_objects }.

    Only keys with `{ selector: ... }` values are object-valued.
    String-valued keys (display text, URLs) are excluded.
    """
    result: dict[str, set[str]] = {}
    _OBJ_KEY_RE = re.compile(r"^\s{2,}(\w+)\s*:\s*\{[^}]*selector:", re.MULTILINE)

    for m_block in _BLOCK_OPEN_RE.finditer(locator_file_content):
        block_name = m_block.group(2)
        body_start = m_block.end()
        depth = 1
        j = body_start
        while j < len(locator_file_content) and depth > 0:
            c = locator_file_content[j]
            if c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
            j += 1
        body = locator_file_content[body_start : j - 1]
        object_keys = {km.group(1) for km in _OBJ_KEY_RE.finditer(body)}
        result[block_name] = object_keys

    return result


def fix_selector_accesses(
    content: str,
    object_keys_by_name: dict[str, set[str]],
) -> tuple[str, list[str]]:
    """
    In a page-object / step file, rewrite `xxxLocators.key` → `xxxLocators.key.selector`
    for every key that is now a structured object `{ selector, intent, stability }`.

    Leaves:
      - string-valued keys unchanged (display text, URLs)
      - function-call entries unchanged: xxxLocators.fn(arg)
      - already-fixed accesses: xxxLocators.key.selector
    """
    changes: list[str] = []
    for loc_name, obj_keys in object_keys_by_name.items():
        if not obj_keys:
            continue
        # Build alternation of all object keys for this locator group
        keys_alt = "|".join(re.escape(k) for k in sorted(obj_keys, key=len, reverse=True))
        pattern = re.compile(
            rf"\b({re.escape(loc_name)})\.({keys_alt})\b"
            r"(?!\.(?:selector|intent|stability))"   # not already .selector/.intent/.stability
            r"(?!\s*\()",                             # not a function call
        )

        def _replace(m: re.Match, ln=loc_name) -> str:
            changes.append(f"Added .selector: {m.group(0)}")
            return f"{m.group(1)}.{m.group(2)}.selector"

        content = pattern.sub(_replace, content)
    return content, changes
