"""Update relative import paths to match Helix-QA directory structure.

The authoritative source for "where does source file X end up in Helix?" is
the mapper. When `_migrate.py` orchestrates the run it builds a lookup of
`{absolute_source_path → helix_target}` and passes it in via `source_to_target`.
The import fixer uses that lookup first — only falling back to heuristic
prediction (`_infer_helix_path`) when the imported file isn't in the mapping
(e.g. the import points to something outside the source tree).

This keeps imports byte-correct after disambiguation, custom mapper rules,
or any other mapper logic the fixer would otherwise have to mirror.
"""
from __future__ import annotations
import re
from pathlib import Path, PurePosixPath


# Known path aliases used in Helix projects
_HELIX_ALIASES = {
    re.compile(r"[Ll]ocator[Hh]ealer"):     "@utils/locators/LocatorHealer",
    re.compile(r"[Ll]ocator[Rr]epository"): "@utils/locators/LocatorRepository",
    re.compile(r"[Tt]iming[Hh]ealer"):      "@utils/locators/TimingHealer",
    re.compile(r"[Vv]isual[Ii]ntent"):      "@utils/locators/VisualIntentChecker",
    re.compile(r"[Hh]ealing[Dd]ashboard"):  "@utils/locators/HealingDashboard",
    re.compile(r"[Bb]ase[Pp]age"):          "./BasePage",
    re.compile(r"pageFixture|PageFixture"):  "@hooks/pageFixture",
}

# Rewrite source-project @alias/... paths to Helix equivalents.
# Locators and page objects that were migrated get their alias updated.
_SOURCE_ALIAS_RE = re.compile(
    r"^(@pages/|@steps/|@features/)(.+?)(?:\.(ts|js))?$"
)


def _rewrite_source_alias(path: str) -> str | None:
    """
    Rewrite a source-project path alias to its Helix equivalent.
    e.g. '@pages/amplify/ui/locators'  →  '@locators/ui-locators.locators'
         '@pages/amplify/ui/login.page' →  '@pages/login.page'
         '@pages/nova/locators'         →  '@locators/nova-locators.locators'
    Returns None if no rewrite is needed or possible.
    """
    m = _SOURCE_ALIAS_RE.match(path)
    if not m:
        return None

    _, tail, _ext = m.groups()
    # `name` is the last path segment, possibly including role suffix like "editModule.page"
    name = Path(tail).name

    from stlc_agents.agent_migration.mapper import _to_kebab, _feature_kebab

    # Locator files (bare "locators" or "novaLocators")
    if re.search(r"[Ll]ocators?$", name):
        base = re.sub(r"[Ll]ocators?$", "", name)
        if not base:
            # Bare "locators" file — use parent directory name as prefix
            parent = Path(tail).parent.name.lower()
            return f"@locators/{parent}-locators.locators"
        return f"@locators/{_to_kebab(base)}-locators.locators"

    # Page objects — name ends with ".page" (no extension) or stem ends with "Page"
    if re.search(r"\.page$", name):
        base = re.sub(r"\.page$", "", name)
        return f"@pages/{_to_kebab(base)}.page"
    if re.search(r"[Pp]age$", Path(name).stem):
        base = re.sub(r"[Pp]age$", "", Path(name).stem)
        return f"@pages/{_to_kebab(base)}.page"

    # Step definitions
    if re.search(r"\.steps?$", name):
        base = re.sub(r"\.steps?$", "", name)
        return f"@steps/{_to_kebab(base)}.steps"
    if re.search(r"[Ss]teps?$", Path(name).stem):
        base = re.sub(r"[Ss]teps?$", "", Path(name).stem)
        return f"@steps/{_to_kebab(base)}.steps"

    # API clients — anything ending in `.api` lives under @helper/api/ in helix-qa
    if re.search(r"\.api$", name):
        return f"@helper/api/{name}"

    return None


def fix_imports(
    content: str,
    source_path: str,
    target_path: str,
    source_root: str,
    source_to_target: dict[str, str] | None = None,
) -> tuple[str, list[str]]:
    """
    Update import paths in content for the new Helix target location.

    source_path       — absolute path of the original file
    target_path       — Helix-relative target path (e.g. 'src/pages/login.page.ts')
    source_root       — absolute path of the source project root
    source_to_target  — authoritative {absolute_source_path → helix_target_rel}
                         lookup from the mapper. When the imported file is
                         present here, its target is used verbatim instead of
                         being predicted by the heuristic. Pass `None` (or omit)
                         to fall back to the legacy guess-only behaviour.
    Returns (updated_content, changes).
    """
    changes: list[str] = []
    lines   = content.split("\n")
    result  = []
    s2t     = source_to_target or {}

    for line in lines:
        new_line = _fix_line(line, source_path, target_path, source_root, s2t, changes)
        result.append(new_line)

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


# ---------------------------------------------------------------------------
# Per-line processing
# ---------------------------------------------------------------------------

_IMPORT_FROM_RE = re.compile(
    r"^(\s*(?:import|export)\s+.*?\s+from\s+)(['\"])(\.[^'\"]+)\2(.*)"
)
# Match absolute alias imports on a single line: import Foo from '@alias/...'
_IMPORT_ALIAS_RE = re.compile(
    r"^(\s*(?:import|export)\s+.*?\s+from\s+)(['\"])(@[^'\"]+)\2(.*)"
)
# Match closing line of a multi-line import: } from '@alias/...'
_IMPORT_ALIAS_CLOSE_RE = re.compile(
    r"^(\s*\}[ \t]+from\s+)(['\"])(@[^'\"]+)\2(.*)"
)
_REQUIRE_RE = re.compile(r"(require\s*\(\s*)(['\"])(\.[^'\"]+)\2(\s*\))")


def _fix_line(
    line: str,
    source_path: str,
    target_path: str,
    source_root: str,
    source_to_target: dict[str, str],
    changes: list[str],
) -> str:
    # ES import/export from '...' (relative paths)
    m = _IMPORT_FROM_RE.match(line)
    if m:
        prefix, quote, rel_path, suffix = m.groups()
        new_path = _remap(rel_path, source_path, target_path, source_root, source_to_target, changes)
        if new_path != rel_path:
            return f"{prefix}{quote}{new_path}{quote}{suffix}"
        return line

    # ES import/export from '@alias/...' (absolute alias paths from source project)
    m = _IMPORT_ALIAS_RE.match(line) or _IMPORT_ALIAS_CLOSE_RE.match(line)
    if m:
        prefix, quote, alias_path, suffix = m.groups()
        new_path = _rewrite_source_alias_via_mapping(
            alias_path, source_root, source_to_target,
        )
        if new_path is None:
            new_path = _rewrite_source_alias(alias_path)
        if new_path and new_path != alias_path:
            changes.append(f"Remapped alias {alias_path!r} → {new_path!r}")
            return f"{prefix}{quote}{new_path}{quote}{suffix}"
        return line

    # CommonJS require('...')
    m = _REQUIRE_RE.search(line)
    if m:
        req_prefix, quote, rel_path, req_suffix = m.groups()
        new_path = _remap(rel_path, source_path, target_path, source_root, source_to_target, changes)
        if new_path != rel_path:
            return line.replace(f"{quote}{rel_path}{quote}", f"{quote}{new_path}{quote}", 1)

    return line


# ---------------------------------------------------------------------------
# Path remapping
# ---------------------------------------------------------------------------

def _remap(
    rel_path: str,
    source_path: str,
    target_path: str,
    source_root: str,
    source_to_target: dict[str, str],
    changes: list[str],
) -> str:
    # Check against known infrastructure aliases first
    for pattern, alias in _HELIX_ALIASES.items():
        if pattern.search(rel_path):
            changes.append(f"Remapped {rel_path!r} → {alias!r} (alias)")
            return alias

    # Resolve the import target relative to the source file
    src_dir = Path(source_path).parent
    raw_target = (src_dir / rel_path).resolve()
    root        = Path(source_root).resolve()

    try:
        rel_to_root = raw_target.relative_to(root).as_posix()
    except ValueError:
        return rel_path  # escapes source root — leave as-is

    # Authoritative path: ask the mapper. We need to match against a file with
    # an extension (the import didn't include one); try `.ts`, `.tsx`, `.js`,
    # `.jsx` and bare-path index.* variants.
    helix_target = _lookup_mapped_target(raw_target, source_to_target)

    # Fallback to heuristic prediction when the imported file isn't a mapped
    # source file (e.g. import of a generated file, or an external module
    # masquerading as a relative path).
    if helix_target is None:
        helix_target = _infer_helix_path(rel_to_root)
    if helix_target is None:
        return rel_path

    new_rel = _relative_helix_import(target_path, helix_target)
    if new_rel != rel_path:
        changes.append(f"Remapped import {rel_path!r} → {new_rel!r}")
    return new_rel


# ---------------------------------------------------------------------------
# Mapper lookup — converts the resolved source path of an imported file into
# the helix target the mapper actually wrote (or will write). The import in
# the source code typically omits the extension, so we probe the common ones.
# ---------------------------------------------------------------------------

_IMPORT_PROBE_SUFFIXES = (".ts", ".tsx", ".js", ".jsx", "")


def _lookup_mapped_target(raw_target: Path, s2t: dict[str, str]) -> str | None:
    """Return the helix target for `raw_target` if the mapper covers it."""
    if not s2t:
        return None
    base = str(raw_target)
    for suffix in _IMPORT_PROBE_SUFFIXES:
        candidate = base + suffix
        if candidate in s2t:
            return _strip_ext(s2t[candidate])
    # The import might point at a directory with an index file inside.
    for index_name in ("index.ts", "index.tsx", "index.js", "index.jsx"):
        candidate = str(raw_target / index_name)
        if candidate in s2t:
            return _strip_ext(s2t[candidate])
    return None


def _strip_ext(target_rel: str) -> str:
    """Drop the file extension from a helix target path for import use."""
    for suffix in (".ts", ".tsx", ".js", ".jsx"):
        if target_rel.endswith(suffix):
            return target_rel[: -len(suffix)]
    return target_rel


# ---------------------------------------------------------------------------
# Alias-via-mapping: resolve a source-project alias (e.g. `@pages/scweb/locators`)
# by reconstructing the absolute source path and looking it up in the mapping.
# Falls back to the legacy heuristic when the alias prefix isn't recognised.
# ---------------------------------------------------------------------------

# Map common source-project alias prefixes to their physical sub-paths inside
# the source tree. Most projects use `@pages/X` for `src/pages/X.ts`, etc.
# Unknown aliases fall through to the heuristic rewriter.
_ALIAS_SUBPATHS = {
    "@pages":    ("src/pages", "pages"),
    "@steps":    ("src/test/steps", "src/steps", "steps"),
    "@features": ("src/test/features", "src/features", "features"),
    "@locators": ("src/locators", "locators"),
    "@helper":   ("src/helper", "src/helpers", "helper", "helpers"),
    "@hooks":    ("src/hooks", "hooks"),
}


def _rewrite_source_alias_via_mapping(
    alias_path: str,
    source_root: str,
    s2t: dict[str, str],
) -> str | None:
    """
    Try to resolve `@alias/sub/path` by probing physical sub-paths in the
    source tree and looking the result up in the mapper's table.
    """
    if not s2t or not alias_path.startswith("@"):
        return None
    head, _, tail = alias_path.partition("/")
    if not tail:
        return None
    subpaths = _ALIAS_SUBPATHS.get(head)
    if not subpaths:
        return None
    root = Path(source_root).resolve()
    for sub in subpaths:
        for suffix in _IMPORT_PROBE_SUFFIXES:
            candidate = root / sub / (tail + suffix)
            mapped = s2t.get(str(candidate.resolve()))
            if mapped:
                return _helix_target_to_alias(_strip_ext(mapped))
        # Also try directory index lookup
        for index_name in ("index.ts", "index.tsx", "index.js", "index.jsx"):
            candidate = root / sub / tail / index_name
            mapped = s2t.get(str(candidate.resolve()))
            if mapped:
                return _helix_target_to_alias(_strip_ext(mapped))
    return None


def _helix_target_to_alias(helix_rel: str) -> str:
    """
    Convert a helix-relative file path (no extension) into the most natural
    `@alias/...` form so callers can import via the configured tsconfig path
    aliases.
    """
    mapping = (
        ("src/locators/",     "@locators/"),
        ("src/pages/",        "@pages/"),
        ("src/test/steps/",   "@steps/"),
        ("src/test/features/","@features/"),
        ("src/hooks/",        "@hooks/"),
        ("src/helper/",       "@helper/"),
        ("src/utils/",        "@utils/"),
    )
    for prefix, alias in mapping:
        if helix_rel.startswith(prefix):
            return alias + helix_rel[len(prefix):]
    return helix_rel


def _infer_helix_path(rel_to_root: str) -> str | None:
    """Guess the Helix file path from a source-root-relative path."""
    name = Path(rel_to_root).stem

    if re.search(r"[Ll]ocators?$", name):
        from stlc_agents.agent_migration.mapper import _to_kebab  # type: ignore[import]
        stem = re.sub(r"[Ll]ocators?$", "", name)
        if not stem:
            # File is literally "locators.ts" — disambiguation would have added
            # the parent directory name as prefix, e.g. ui/ → ui-locators.locators
            parent = Path(rel_to_root).parent.name.lower()
            return f"src/locators/{parent}-locators.locators"
        return f"src/locators/{_to_kebab(stem)}.locators"

    # Page-object import patterns:
    #   `Xxx.page`     (stem = "Xxx.page", or "Xxx" after one strip)
    #   `XxxPage`      (PascalCase suffix)
    rel_path = rel_to_root
    if rel_path.endswith(".page") or re.search(r"\.page(\.[jt]s)?$", rel_path):
        from stlc_agents.agent_migration.mapper import _to_kebab  # type: ignore[import]
        # Get the bare class name: ".../Foo.page" → "Foo"
        base = Path(rel_path).name
        base = re.sub(r"\.[jt]s$", "", base)
        base = re.sub(r"\.page$", "", base, flags=re.IGNORECASE)
        return f"src/pages/{_to_kebab(base)}.page"
    if re.search(r"[Pp]age$", name):
        from stlc_agents.agent_migration.mapper import _to_kebab  # type: ignore[import]
        stem = re.sub(r"[Pp]age$", "", name)
        return f"src/pages/{_to_kebab(stem)}.page"

    if rel_path.endswith(".steps") or re.search(r"\.steps?(\.[jt]s)?$", rel_path):
        from stlc_agents.agent_migration.mapper import _to_kebab  # type: ignore[import]
        base = Path(rel_path).name
        base = re.sub(r"\.[jt]s$", "", base)
        base = re.sub(r"\.steps?$", "", base, flags=re.IGNORECASE)
        return f"src/test/steps/{_to_kebab(base)}.steps"
    if re.search(r"[Ss]teps?$", name):
        from stlc_agents.agent_migration.mapper import _to_kebab  # type: ignore[import]
        stem = re.sub(r"[Ss]teps?$", "", name)
        return f"src/test/steps/{_to_kebab(stem)}.steps"

    # Helper / util files (anywhere under `helpers/` or `utils/`). The mapper
    # routes these to `src/helper/<preserved-subpath>` — match its behavior so
    # imports get rewritten to the new flat location instead of the legacy
    # source-relative path.
    m = re.match(r"^(?:helpers?|utils?)/(.+)$", rel_to_root)
    if m:
        sub = m.group(1).rsplit(".", 1)[0]  # drop extension
        return f"src/helper/{sub}"

    # Fallback for files at the source root (no directory prefix). Things like
    # `global-cleanup.ts`, `global-variableSetup.ts` are raw-copied to the
    # helix root and keep their original name — recomputing the relative path
    # against the new file location is still needed.
    if "/" not in rel_to_root:
        return rel_to_root.rsplit(".", 1)[0]   # drop extension; helix-root-relative

    return None


def _relative_helix_import(from_helix: str, to_helix: str) -> str:
    """Compute a relative import path between two Helix-relative paths."""
    if to_helix.startswith("@"):
        return to_helix

    from_dir = PurePosixPath(from_helix).parent
    to_path  = PurePosixPath(to_helix)

    try:
        rel = to_path.relative_to(from_dir)
        return "./" + str(rel)
    except ValueError:
        pass

    # Walk up from from_dir until we find the common ancestor
    from_parts = from_dir.parts
    to_parts   = to_path.parts

    common = 0
    for f, t in zip(from_parts, to_parts):
        if f == t:
            common += 1
        else:
            break

    up   = len(from_parts) - common
    down = to_parts[common:]
    return "../" * up + "/".join(down)
