"""Cross-reference index for codebase symbol analysis.

Builds a bidirectional index of symbol definitions and references,
file-level import dependencies, and dependents from parsed AST data.
"""

from __future__ import annotations

import logging
import re
import threading
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from attocode_intel._internal.integrations.context.index_store import IndexStore

logger = logging.getLogger(__name__)


def _split_name_tokens(name: str) -> list[str]:
    """Split a symbol name into lowercase tokens.

    Handles camelCase, PascalCase, snake_case, UPPER_CASE, and mixtures.

    Examples::

        "parseConfig"  -> ["parse", "config"]
        "parse_config" -> ["parse", "config"]
        "HTTPServer"   -> ["http", "server"]
        "HTMLParser"   -> ["html", "parser"]
        "_main"        -> ["main"]
        "IO"           -> ["io"]
    """
    # First split on underscores
    parts = name.split("_")
    tokens: list[str] = []
    for part in parts:
        if not part:
            continue
        # Split camelCase / PascalCase: insert boundary before uppercase
        # that follows lowercase, or before uppercase followed by lowercase
        # (to handle "HTTPServer" -> "HTTP", "Server")
        sub = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", part)
        sub = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", sub)
        for s in sub.split("_"):
            if s:
                tokens.append(s.lower())
    return tokens


@dataclass(slots=True)
class SymbolRef:
    """A reference (call site, import, or attribute access) to a symbol."""

    symbol_name: str        # "parse_file" or "CodeIndex.build"
    ref_kind: str           # "call" | "import" | "attribute"
    file_path: str
    line: int
    source: str = "tree-sitter"  # "tree-sitter" | "lsp"
    syntax_name: str = ""  # identifier/member expression as written in source
    # Qualified name of the function/method enclosing this reference. Empty
    # when the reference is at module top level or the enclosing symbol
    # could not be resolved. Populated for call-graph construction.
    caller_qualified_name: str = ""
    definition_key: str = ""  # JSON [definition file, start line, qualified name] from a selected LSP query


@dataclass(slots=True)
class SymbolLocation:
    """Location where a symbol is defined."""

    name: str
    qualified_name: str     # "CodebaseContextManager.discover_files"
    kind: str               # "function" | "class" | "method"
    file_path: str
    start_line: int
    end_line: int
    source: str = "tree-sitter"  # "tree-sitter" | "lsp"


@dataclass(slots=True)
class CrossRefIndex:
    """Bidirectional cross-reference index for a codebase.

    Tracks symbol definitions, call sites, file-level import dependencies,
    and reverse dependents.  Includes an inverted name index for fast
    multi-strategy symbol search.

    Optionally backed by an ``IndexStore`` for persistence across restarts.
    """

    # qualified_name -> list of definition locations
    definitions: dict[str, list[SymbolLocation]] = field(default_factory=dict)
    # symbol name -> list of reference sites
    references: dict[str, list[SymbolRef]] = field(default_factory=dict)
    # file path -> set of qualified symbol names defined in that file
    file_symbols: dict[str, set[str]] = field(default_factory=dict)
    # References to the same locations, grouped by file and qualified name.
    # File navigation must not scan every workspace definition of common names.
    _file_definitions: dict[str, dict[str, list[SymbolLocation]]] = field(
        default_factory=dict, init=False, repr=False, compare=False,
    )
    # file path -> set of file paths it imports from
    file_dependencies: dict[str, set[str]] = field(default_factory=dict)
    # file path -> set of file paths that import it
    file_dependents: dict[str, set[str]] = field(default_factory=dict)
    # caller qualified name -> set of callee bare/qualified names
    call_edges: dict[str, set[str]] = field(default_factory=dict)
    # callee bare/qualified name -> set of caller qualified names (reverse)
    callers_of: dict[str, set[str]] = field(default_factory=dict)

    # --- Inverted name indexes (populated by add_definition) ---
    # bare name -> set of qualified names
    _name_to_qnames: dict[str, set[str]] = field(default_factory=dict)
    _known_names_snapshot: frozenset[str] | None = field(default=None, repr=False, compare=False)
    # lowercased bare name -> set of qualified names
    _lower_to_qnames: dict[str, set[str]] = field(default_factory=dict)
    # individual token -> set of qualified names
    _tokens_to_qnames: dict[str, set[str]] = field(default_factory=dict)

    # Optional persistent store (set via set_store)
    _store: Any = field(default=None, repr=False)

    # Re-entrant lock guarding all mutating methods plus the iterating
    # reads that need a consistent snapshot. Reentrancy matters because
    # ``add_reference`` may call ``add_call_edge`` and ``remove_file``
    # walks every other map. Background hydration (``ASTService``'s
    # daemon thread) and LSP/MCP request threads can mutate the index
    # concurrently, so the lock prevents corrupted dicts and torn lists.
    _lock: threading.RLock = field(
        default_factory=threading.RLock, repr=False, compare=False,
    )

    def __post_init__(self) -> None:
        for locations in self.definitions.values():
            for loc in locations:
                self._file_definitions.setdefault(loc.file_path, {}).setdefault(loc.qualified_name, []).append(loc)
                self.file_symbols.setdefault(loc.file_path, set()).add(loc.qualified_name)

    def get_file_definitions(self, file_path: str) -> list[SymbolLocation]:
        """Return only this file's definitions, preserving within-name order."""
        with self._lock:
            grouped = self._file_definitions.get(file_path, {})
            return [loc for name in self.file_symbols.get(file_path, ())
                    for loc in grouped.get(name, ())]

    def set_store(self, store: IndexStore) -> None:
        """Attach an IndexStore for write-through persistence."""
        self._store = store

    def persist_file(self, file_path: str) -> None:
        """Write-through: save current in-memory symbols/refs for a file to the store."""
        if self._store is None:
            return
        with self._lock:
            self._persist_file_locked(file_path)

    def persist_files(self, file_paths: set[str]) -> None:
        """Persist a changed batch without rescanning all references per file."""
        if self._store is None:
            return
        with self._lock:
            grouped: dict[str, list[SymbolRef]] = {}
            for refs in self.references.values():
                for ref in refs:
                    if ref.file_path in file_paths:
                        grouped.setdefault(ref.file_path, []).append(ref)
            paths = sorted(file_paths)
            for offset in range(0, len(paths), 100):
                self._store.save_analysis_batch([
                    (path, *self._file_evidence_locked(path, grouped.get(path, [])))
                    for path in paths[offset:offset + 100]
                ])

    def _persist_file_locked(self, file_path: str, references=None) -> None:
        if self._store is None:
            return
        self._store.save_analysis_batch([(file_path, *self._file_evidence_locked(file_path, references))])

    def _file_evidence_locked(self, file_path: str, references=None):
        # Collect symbols for this file
        sym_dicts: list[dict[str, Any]] = []
        for loc in self.get_file_definitions(file_path):
            sym_dicts.append({
                "name": loc.name,
                "qualified_name": loc.qualified_name,
                "kind": loc.kind,
                "line": loc.start_line,
                "end_line": loc.end_line,
                "source": loc.source,
            })

        # Collect references originating from this file
        ref_dicts: list[dict[str, Any]] = []
        refs = references if references is not None else (
            ref for refs in self.references.values() for ref in refs if ref.file_path == file_path
        )
        for ref in refs:
            ref_dicts.append({
                "symbol_name": ref.symbol_name,
                "ref_kind": ref.ref_kind,
                "line": ref.line,
                "column": 0,
                "source": ref.source,
                "syntax_name": ref.syntax_name,
                "caller_qualified_name": ref.caller_qualified_name,
                "definition_key": ref.definition_key,
            })
        return sym_dicts, ref_dicts

    def load_from_store(self) -> int:
        """Bulk-load symbols, references, and dependencies from the store.

        Returns the number of files loaded.
        """
        if self._store is None:
            return 0
        with self._lock:
            return self._load_from_store_locked()

    def _load_from_store_locked(self) -> int:
        if self._store is None:
            return 0

        stored_symbols = self._store.load_symbols()
        stored_refs = self._store.load_references()
        stored_deps = self._store.load_dependencies()

        files_loaded: set[str] = set()

        for s in stored_symbols:
            loc = SymbolLocation(
                name=s.name,
                qualified_name=s.qualified_name,
                kind=s.kind,
                file_path=s.file_path,
                start_line=s.line,
                end_line=s.end_line,
                source=s.source,
            )
            self.add_definition(loc)
            files_loaded.add(s.file_path)

        for r in stored_refs:
            ref = SymbolRef(
                symbol_name=r.symbol_name,
                ref_kind=r.ref_kind,
                file_path=r.file_path,
                line=r.line,
                source=r.source,
                syntax_name=getattr(r, "syntax_name", ""),
                caller_qualified_name=getattr(r, "caller_qualified_name", ""),
                definition_key=getattr(r, "definition_key", ""),
            )
            self.add_reference(ref)

        for src, targets in stored_deps.items():
            for tgt in targets:
                self.add_file_dependency(src, tgt)

        logger.debug(
            "CrossRefIndex loaded from store: %d files, %d symbols, %d refs",
            len(files_loaded), len(stored_symbols), len(stored_refs),
        )
        return len(files_loaded)

    def merge_lsp_results(
        self,
        file_path: str,
        definitions: list[SymbolLocation],
        references: list[SymbolRef],
        *, verified_symbol: bool = False,
    ) -> int:
        """Merge LSP-sourced results into the index.

        Deduplicates by (qualified_name, file, line) — LSP wins on conflict.
        Returns the number of new entries added.
        """
        with self._lock:
            return self._merge_lsp_results_locked(file_path, definitions, references, verified_symbol=verified_symbol)

    def _merge_lsp_results_locked(
        self,
        file_path: str,
        definitions: list[SymbolLocation],
        references: list[SymbolRef],
        *, verified_symbol: bool = False,
    ) -> int:
        added = 0

        for loc in definitions:
            # Create a copy with source="lsp" — don't mutate caller's input
            lsp_loc = SymbolLocation(
                name=loc.name, qualified_name=loc.qualified_name,
                kind=loc.kind, file_path=loc.file_path,
                start_line=loc.start_line, end_line=loc.end_line,
                source="lsp",
            )
            existing = self.definitions.get(lsp_loc.qualified_name, [])
            # Check for duplicate at same location
            dup = False
            for i, ex in enumerate(existing):
                if ex.file_path == lsp_loc.file_path and ex.start_line == lsp_loc.start_line:
                    # LSP wins: replace tree-sitter entry
                    existing[i] = lsp_loc
                    local = self._file_definitions[lsp_loc.file_path][lsp_loc.qualified_name]
                    for j, prior in enumerate(local):
                        if prior is ex:
                            local[j] = lsp_loc
                            break
                    dup = True
                    break
            if not dup:
                self.add_definition(lsp_loc)
                added += 1

        for ref in references:
            lsp_ref = SymbolRef(
                symbol_name=ref.symbol_name, ref_kind=ref.ref_kind,
                file_path=ref.file_path, line=ref.line,
                source="lsp",
                syntax_name=ref.syntax_name,
                caller_qualified_name=ref.caller_qualified_name,
                definition_key=ref.definition_key,
            )
            # LSP wins on (file, line) collision — but ONLY when the
            # incoming ref carries a resolved caller. Without one, the
            # legacy ``ingest_lsp_results`` fallback path stores the
            # *enclosing function* in ``symbol_name`` (not the callee),
            # so replacing a tree-sitter entry would invert the edge.
            # In that case we dedupe-skip, preserving the existing
            # higher-precision tree-sitter ref.
            existing = self.references.get(lsp_ref.symbol_name, [])
            dup_index: int | None = None
            for i, ex in enumerate(existing):
                if (ex.file_path == lsp_ref.file_path and ex.line == lsp_ref.line
                        and (not ex.definition_key or ex.definition_key == lsp_ref.definition_key)):
                    dup_index = i
                    break

            if dup_index is not None:
                if lsp_ref.caller_qualified_name or verified_symbol:
                    if not lsp_ref.syntax_name:
                        lsp_ref.syntax_name = existing[dup_index].syntax_name
                    existing[dup_index] = lsp_ref
                    if lsp_ref.ref_kind == "call":
                        self.add_call_edge(
                            lsp_ref.caller_qualified_name,
                            lsp_ref.symbol_name,
                        )
                # else: no caller info — leave the tree-sitter entry alone.
            else:
                self.add_reference(lsp_ref)
                added += 1

        # Persist if store is available
        self.persist_file(file_path)
        return added

    def known_symbol_names(self) -> frozenset[str]:
        """Share one immutable name snapshot until definitions add/remove a name.

        Bulk reference indexing must not copy every repository name per file.
        Existing readers retain a consistent snapshot while edits invalidate the
        next lookup under the same lock as the inverted name index.
        """
        with self._lock:
            if self._known_names_snapshot is None:
                self._known_names_snapshot = frozenset(self._name_to_qnames)
            return self._known_names_snapshot

    def add_definition(self, loc: SymbolLocation) -> None:
        """Register a symbol definition."""
        with self._lock:
            self.definitions.setdefault(loc.qualified_name, []).append(loc)
            self.file_symbols.setdefault(loc.file_path, set()).add(loc.qualified_name)
            self._file_definitions.setdefault(loc.file_path, {}).setdefault(loc.qualified_name, []).append(loc)
            # Populate inverted indexes
            bare = loc.qualified_name.rsplit(".", 1)[-1]
            if bare not in self._name_to_qnames:
                self._known_names_snapshot = None
            self._name_to_qnames.setdefault(bare, set()).add(loc.qualified_name)
            self._lower_to_qnames.setdefault(bare.lower(), set()).add(loc.qualified_name)
            for token in _split_name_tokens(bare):
                self._tokens_to_qnames.setdefault(token, set()).add(loc.qualified_name)

    def add_reference(self, ref: SymbolRef) -> None:
        """Register a symbol reference."""
        with self._lock:
            self.references.setdefault(ref.symbol_name, []).append(ref)
            if ref.ref_kind == "call" and ref.caller_qualified_name:
                self.add_call_edge(ref.caller_qualified_name, ref.symbol_name)

    def add_call_edge(self, caller: str, callee: str) -> None:
        """Record that ``caller`` calls ``callee``.

        ``caller`` should be a qualified name (e.g. ``MyClass.do_thing`` or a
        bare top-level function name). ``callee`` is whatever the source
        wrote — bare name when type/scope resolution is unavailable; the
        consumer can disambiguate against ``self.definitions``.
        """
        if not caller or not callee:
            return
        with self._lock:
            self.call_edges.setdefault(caller, set()).add(callee)
            self.callers_of.setdefault(callee, set()).add(caller)

    def add_file_dependency(self, source: str, target: str) -> None:
        """Record that *source* imports from *target*."""
        with self._lock:
            self.file_dependencies.setdefault(source, set()).add(target)
            self.file_dependents.setdefault(target, set()).add(source)

    def remove_file(self, file_path: str) -> None:
        """Remove all index entries for a file (in-memory and store)."""
        with self._lock:
            self._remove_file_locked(file_path)

    def _remove_file_locked(self, file_path: str) -> None:
        if self._store is not None:
            self._store.remove_file(file_path)
        # Snapshot callers defined here BEFORE we mutate file_symbols below.
        callers_in_file = set(self.file_symbols.get(file_path, set()))
        # Remove definitions and clean up inverted indexes
        for qname in list(self.file_symbols.get(file_path, [])):
            defs = self.definitions.get(qname, [])
            self.definitions[qname] = [d for d in defs if d.file_path != file_path]
            if not self.definitions[qname]:
                del self.definitions[qname]
                # Clean up inverted indexes — qname fully gone
                bare = qname.rsplit(".", 1)[-1]
                s = self._name_to_qnames.get(bare)
                if s:
                    s.discard(qname)
                    if not s:
                        del self._name_to_qnames[bare]
                        self._known_names_snapshot = None
                s = self._lower_to_qnames.get(bare.lower())
                if s:
                    s.discard(qname)
                    if not s:
                        del self._lower_to_qnames[bare.lower()]
                for token in _split_name_tokens(bare):
                    s = self._tokens_to_qnames.get(token)
                    if s:
                        s.discard(qname)
                        if not s:
                            del self._tokens_to_qnames[token]
        self.file_symbols.pop(file_path, None)
        self._file_definitions.pop(file_path, None)

        # Remove references originating from this file
        for name, refs in list(self.references.items()):
            self.references[name] = [r for r in refs if r.file_path != file_path]
            if not self.references[name]:
                del self.references[name]

        # Drop call edges whose caller was defined in this file. We can't
        # reliably attribute callees from the file path alone (callees are
        # bare names), so this is a best-effort scrub keyed on caller.
        for caller in callers_in_file:
            for callee in self.call_edges.pop(caller, set()):
                rev = self.callers_of.get(callee)
                if rev is not None:
                    rev.discard(caller)
                    if not rev:
                        del self.callers_of[callee]

        # Remove dependency edges
        for dep in self.file_dependencies.pop(file_path, set()):
            dependents = self.file_dependents.get(dep)
            if dependents:
                dependents.discard(file_path)
                if not dependents:
                    del self.file_dependents[dep]

        # Remove as a dependent target
        for src in list(self.file_dependents.get(file_path, [])):
            deps = self.file_dependencies.get(src)
            if deps:
                deps.discard(file_path)
                if not deps:
                    del self.file_dependencies[src]
        self.file_dependents.pop(file_path, None)

    # ------------------------------------------------------------------
    # Search
    # ------------------------------------------------------------------

    def get_definitions(self, symbol_name: str, *, include_shadowed: bool = False) -> list[SymbolLocation]:
        """Look up definitions for a symbol (exact or suffix match)."""
        # The lock keeps a background index writer from changing the dict during the scan.
        with self._lock:
            # Try exact qualified name first
            if symbol_name in self.definitions and not include_shadowed:
                return self.definitions[symbol_name]
            # Try suffix match (e.g. "discover_files" matches "CodebaseContextManager.discover_files")
            results: list[SymbolLocation] = []
            for qname, locs in self.definitions.items():
                if qname == symbol_name or qname.endswith(f".{symbol_name}"):
                    results.extend(locs)
            return results

    def search_definitions(
        self,
        symbol_name: str,
        *,
        limit: int = 50,
        kind_filter: str = "",
    ) -> list[tuple[SymbolLocation, float]]:
        """Multi-strategy symbol search returning (location, score) pairs."""
        with self._lock:
            return self._search_definitions_locked(symbol_name, limit=limit, kind_filter=kind_filter)

    def _search_definitions_locked(
        self,
        symbol_name: str,
        *,
        limit: int,
        kind_filter: str,
    ) -> list[tuple[SymbolLocation, float]]:
        """Search the definitions. The caller holds the lock.

        Search strategies (in priority order):
        1. Exact qualified name match (score 1.0)
        2. Exact bare name match via index (score 0.95)
        3. Case-insensitive bare name match (score 0.85)
        4. Prefix match on bare name (score 0.75)
        5. Substring match on bare name (score 0.60)
        6. Token overlap (camelCase/snake_case) (score 0.50)
        """
        # Collect candidates as {qualified_name: best_match_score}
        candidates: dict[str, float] = {}

        # 1. Exact qualified name
        if symbol_name in self.definitions:
            candidates[symbol_name] = 1.0

        # 2. Exact bare name via inverted index — O(1)
        bare_qnames = self._name_to_qnames.get(symbol_name, set())
        for qn in bare_qnames:
            if qn not in candidates:
                candidates[qn] = 0.95

        # 3. Case-insensitive bare name — O(1)
        lower_qnames = self._lower_to_qnames.get(symbol_name.lower(), set())
        for qn in lower_qnames:
            if qn not in candidates:
                candidates[qn] = 0.85

        name_lower = symbol_name.lower()

        # 4. Prefix match on bare names — scan index keys
        for bare, qnames in self._name_to_qnames.items():
            if bare.lower().startswith(name_lower) and len(bare) > len(symbol_name):
                for qn in qnames:
                    if qn not in candidates:
                        candidates[qn] = 0.75

        # 5. Substring match on bare names
        if len(symbol_name) >= 2:
            for bare, qnames in self._name_to_qnames.items():
                if name_lower in bare.lower() and bare.lower() != name_lower:
                    for qn in qnames:
                        if qn not in candidates:
                            candidates[qn] = 0.60

        # 6. Token overlap (camelCase/snake_case)
        query_tokens = set(_split_name_tokens(symbol_name))
        if query_tokens:
            token_hits: dict[str, int] = {}
            for token in query_tokens:
                for qn in self._tokens_to_qnames.get(token, set()):
                    if qn not in candidates:
                        token_hits[qn] = token_hits.get(qn, 0) + 1
            for qn, hit_count in token_hits.items():
                # Score proportional to fraction of query tokens matched
                score = 0.50 * (hit_count / len(query_tokens))
                if score >= 0.25:  # at least half the tokens must match
                    candidates[qn] = score

        if not candidates:
            return []

        # Resolve candidates to SymbolLocations and compute composite scores
        scored: list[tuple[SymbolLocation, float]] = []
        for qname, match_score in candidates.items():
            locs = self.definitions.get(qname, [])
            for loc in locs:
                if kind_filter and loc.kind != kind_filter:
                    continue
                score = _rank_score(loc, match_score)
                scored.append((loc, score))

        # Cold indexing, persisted restoration and name-index sets can enumerate
        # equal matches differently. Break ties before the bounded page is cut.
        scored.sort(key=lambda x: (
            -x[1], x[0].file_path, x[0].start_line, x[0].end_line,
            x[0].qualified_name, x[0].name, x[0].kind, x[0].source,
        ))
        return scored[:limit]

    def get_references(self, symbol_name: str) -> list[SymbolRef]:
        """Look up all call sites / references for a symbol (exact or suffix match)."""
        # Include exact and qualified candidates; an exact key must not hide others.
        # Suffix match (e.g. "clear" matches refs keyed as "MCPMetaTools.clear")
        results: list[SymbolRef] = []
        with self._lock:
            for ref_name, refs in self.references.items():
                if ref_name == symbol_name or ref_name.endswith(f".{symbol_name}"):
                    results.extend(refs)
        return results

    def get_dependents(self, file_path: str) -> set[str]:
        """Files that import from *file_path*."""
        return self.file_dependents.get(file_path, set())

    def get_dependencies(self, file_path: str) -> set[str]:
        """Files that *file_path* imports from."""
        return self.file_dependencies.get(file_path, set())

    # ------------------------------------------------------------------
    # Call graph
    # ------------------------------------------------------------------

    def get_callees(self, caller: str, *, depth: int = 1) -> set[str]:
        """Return symbols called by ``caller`` up to ``depth`` hops.

        depth=1 returns direct callees only; depth=N follows up to N hops
        through the in-memory ``call_edges`` map. Cycles are handled by
        an internal visited set so the traversal always terminates.
        """
        if depth < 1:
            return set()
        with self._lock:
            visited: set[str] = set()
            frontier: set[str] = {caller}
            result: set[str] = set()
            for _ in range(depth):
                next_frontier: set[str] = set()
                for node in frontier:
                    if node in visited:
                        continue
                    visited.add(node)
                    callees = self.call_edges.get(node, set())
                    result.update(callees)
                    next_frontier.update(callees - visited)
                if not next_frontier:
                    break
                frontier = next_frontier
            return result

    def get_callers(self, callee: str, *, depth: int = 1) -> set[str]:
        """Return symbols that call ``callee`` up to ``depth`` hops back."""
        if depth < 1:
            return set()
        with self._lock:
            visited: set[str] = set()
            frontier: set[str] = {callee}
            result: set[str] = set()
            for _ in range(depth):
                next_frontier: set[str] = set()
                for node in frontier:
                    if node in visited:
                        continue
                    visited.add(node)
                    callers = self.callers_of.get(node, set())
                    result.update(callers)
                    next_frontier.update(callers - visited)
                if not next_frontier:
                    break
                frontier = next_frontier
            return result


def _rank_score(loc: SymbolLocation, match_score: float) -> float:
    """Composite ranking combining match quality and symbol importance."""
    importance = 0.0
    if not loc.name.startswith("_"):
        importance += 0.02
    kind_boost = {"class": 0.03, "interface": 0.03, "function": 0.01, "method": 0.0}
    importance += kind_boost.get(loc.kind, 0.0)
    # LSP-sourced entries are more precise
    if loc.source == "lsp":
        importance += 0.05
    return match_score + importance
