require "json"
require "./file_index"
require "../workspace_analyzer"

module Lattice
  module Index
    # WorkspaceIndex aggregates `FileIndex` entries for every source file in the
    # workspace and provides workspace-wide symbol lookup used by hover,
    # completion, go-to-definition, and find-references handlers.
    #
    # Persistence strategy:
    # - On each file save / indexing run, the updated `FileIndex` is written to
    #   `.lattice-cache/index/<sha256-of-path>.json`.
    # - On startup the cache directory is scanned and valid (non-stale) entries
    #   are loaded, avoiding a cold full-workspace scan for large projects.
    class WorkspaceIndex
      CACHE_SUBDIR = "index"

      @entries = Hash(String, FileIndex).new  # keyed by file path
      @cache_dir : String?

      def initialize(workspace_root : String? = nil)
        if root = workspace_root
          dir = File.join(root, ".lattice-cache", CACHE_SUBDIR)
          Dir.mkdir_p(dir)
          @cache_dir = dir
          load_cache(dir)
        end
      end

      # ---------------------------------------------------------------------------
      # Index mutation
      # ---------------------------------------------------------------------------

      # Index a single source file from *source* text.  Skips re-indexing if the
      # stored entry is still fresh.  Returns the (possibly cached) `FileIndex`.
      def index_file(path : String, source : String) : FileIndex
        key = normalize_key(path)
        hash = Digest::SHA256.hexdigest(source)[0, 8].to_u32(16)

        if existing = @entries[key]?
          return existing unless existing.stale?(source)
        end

        symbols = extract_symbols(path, source)
        entry = FileIndex.new(path, hash, symbols)
        @entries[key] = entry
        persist(entry)
        entry
      end

      # Remove a file from the index (called on `textDocument/didClose` or file
      # deletion).
      def remove(path : String)
        @entries.delete(normalize_key(path))
        if dir = @cache_dir
          File.delete(cache_path(dir, path)) rescue nil
        end
      end

      # ---------------------------------------------------------------------------
      # Lookup helpers
      # ---------------------------------------------------------------------------

      # Find every `SymbolEntry` whose name exactly matches *name* across all
      # indexed files.
      def find_by_name(name : String) : Array(SymbolEntry)
        results = [] of SymbolEntry
        @entries.each_value { |fi| results.concat(fi.find_by_name(name)) }
        results
      end

      # Find every `SymbolEntry` whose name exactly matches *name*, or whose
      # fully-qualified name ends with "::name". Methods and nested types are
      # indexed under their fully-qualified name (e.g. "MyClass::run"), but a
      # caller resolving the identifier under the cursor usually only has the
      # bare short name ("run") to search with.
      def find_by_name_or_suffix(name : String) : Array(SymbolEntry)
        results = [] of SymbolEntry
        suffix = "::#{name}"
        @entries.each_value do |fi|
          fi.symbols.each do |s|
            results << s if s.name == name || s.name.ends_with?(suffix)
          end
        end
        results
      end

      # Find every `SymbolEntry` in the given *file_uri*.
      def symbols_for_file(file_uri : String) : Array(SymbolEntry)
        path = uri_to_path(file_uri)
        @entries[normalize_key(path)]?.try(&.symbols) || [] of SymbolEntry
      end

      # Return all symbols of a given *kind* across the whole workspace.
      def all_of_kind(kind : String) : Array(SymbolEntry)
        results = [] of SymbolEntry
        @entries.each_value { |fi| results.concat(fi.symbols_of_kind(kind)) }
        results
      end

      # Fuzzy / substring search across all symbol names; used for
      # `workspace/symbol`.
      def search(query : String) : Array(SymbolEntry)
        q = query.downcase
        results = [] of SymbolEntry
        @entries.each_value do |fi|
          fi.symbols.each { |s| results << s if s.name.downcase.includes?(q) }
        end
        results
      end

      # ---------------------------------------------------------------------------
      # Private helpers
      # ---------------------------------------------------------------------------

      # Build a `SymbolEntry` array by regex-scanning *source*.  This is the
      # same lightweight approach used by `WorkspaceAnalyzer` — a full PSI walk
      # is added in Session 4 when the PSI layer is complete.
      private def extract_symbols(path : String, source : String) : Array(SymbolEntry)
        symbols = [] of SymbolEntry
        lines = source.split('\n')
        namespace_stack = [] of String

        lines.each_with_index do |line, i|
          doc = extract_doc(lines, i)

          if m = line.match(/^\s*class\s+([\w:]+)(?:\s*<\s*([\w:]+))?/)
            name = m[1]
            parent = m[2]? || "Object"
            full = (namespace_stack + [name]).join("::")
            symbols << SymbolEntry.new(full, "class", parent, path, i, name_col(m, line), line.strip, doc)
            namespace_stack << name
          elsif m = line.match(/^\s*module\s+([\w:]+)/)
            name = m[1]
            full = (namespace_stack + [name]).join("::")
            symbols << SymbolEntry.new(full, "module", "Module", path, i, name_col(m, line), line.strip, doc)
            namespace_stack << name
          elsif m = line.match(/^\s*struct\s+([\w:]+)/)
            name = m[1]
            full = (namespace_stack + [name]).join("::")
            symbols << SymbolEntry.new(full, "struct", "Struct", path, i, name_col(m, line), line.strip, doc)
            namespace_stack << name
          elsif m = line.match(/^\s*enum\s+([\w:]+)/)
            name = m[1]
            full = (namespace_stack + [name]).join("::")
            symbols << SymbolEntry.new(full, "enum", "Enum", path, i, name_col(m, line), line.strip, doc)
            namespace_stack << name
          elsif m = line.match(/^\s*lib\s+([\w:]+)/)
            name = m[1]
            full = (namespace_stack + [name]).join("::")
            symbols << SymbolEntry.new(full, "lib", "Lib", path, i, name_col(m, line), line.strip, doc)
            namespace_stack << name
          elsif m = line.match(/^\s*(?:private\s+)?def\s+(?:self\.)?(\w[\w?!]*)\s*(?:\([^)]*\))?\s*(?::\s*(\w[\w?!]*))?/)
            name = m[1]
            ret  = m[2]? || "Void"
            full = namespace_stack.empty? ? name : "#{namespace_stack.join("::")}::#{name}"
            symbols << SymbolEntry.new(full, "method", ret, path, i, name_col(m, line), line.strip, doc)
          elsif m = line.match(/^\s*(?:property|getter|setter)\s+(\w+)\s*:\s*([\w?!]+)/)
            prop = m[1]
            type = m[2]
            full = namespace_stack.empty? ? "@#{prop}" : "#{namespace_stack.join("::")}::@#{prop}"
            symbols << SymbolEntry.new(full, "property", type, path, i, name_col(m, line), line.strip, doc)
          elsif m = line.match(/^\s*@(\w+)\s*:\s*([\w?!]+)/)
            var  = "@#{m[1]}"
            type = m[2]
            full = namespace_stack.empty? ? var : "#{namespace_stack.join("::")}::#{var}"
            symbols << SymbolEntry.new(full, "instance_variable", type, path, i, name_col(m, line), line.strip, doc)
          elsif m = line.match(/^\s*([A-Z][A-Za-z_0-9]*)\s*=\s*(.+)/)
            name = m[1]
            full = (namespace_stack + [name]).join("::")
            symbols << SymbolEntry.new(full, "constant", "Constant", path, i, name_col(m, line), line.strip, doc)
          elsif line.match(/^\s*end\s*$/)
            namespace_stack.pop unless namespace_stack.empty?
          end
        end

        symbols
      end

      private def extract_doc(lines : Array(String), line_num : Int32) : String?
        docs = [] of String
        cur = line_num - 1
        while cur >= 0
          l = lines[cur].strip
          if l.starts_with?("#")
            docs.unshift(l.sub(/^#\s?/, ""))
            cur -= 1
          elsif l.empty?
            cur -= 1
          else
            break
          end
        end
        docs.empty? ? nil : docs.join("\n")
      end

      private def col(line : String) : Int32
        line.size - line.lstrip.size
      end

      # Return the column of the symbol *name* itself (capture group 1),
      # falling back to the line's indentation width only if the group
      # somehow did not participate in the match. Previously every call site
      # used `col(line)` (indentation width) directly, which made call/type
      # hierarchy items point at the start of the declaration line instead of
      # the name token.
      private def name_col(m : Regex::MatchData, line : String) : Int32
        m.begin(1) || col(line)
      end

      # ---------------------------------------------------------------------------
      # Cache persistence
      # ---------------------------------------------------------------------------

      private def persist(entry : FileIndex)
        dir = @cache_dir
        return unless dir
        File.write(cache_path(dir, entry.path), entry.to_json)
      rescue ex
        STDERR.puts "(Lattice) WorkspaceIndex: failed to persist #{entry.path}: #{ex.message}"
      end

      private def load_cache(dir : String)
        return unless Dir.exists?(dir)
        Dir.each_child(dir) do |f|
          next unless f.ends_with?(".json")
          full = File.join(dir, f)
          json = File.read(full)
          entry = FileIndex.from_json(json)
          # Only load the entry if the file still exists on disk.
          if File.exists?(entry.path)
            source = File.read(entry.path)
            @entries[normalize_key(entry.path)] = entry unless entry.stale?(source)
          end
        rescue
          next
        end
      rescue ex
        STDERR.puts "(Lattice) WorkspaceIndex: cache load error: #{ex.message}"
      end

      private def cache_path(dir : String, file_path : String) : String
        hex = Digest::SHA256.hexdigest(file_path)[0, 16]
        File.join(dir, "#{hex}.json")
      end

      private def uri_to_path(uri : String) : String
        return uri unless uri.starts_with?("file://")
        p = uri.sub(/^file:\/\//, "")
        p = URI.decode(p)
        p = p[1..] if p =~ /^\/([a-zA-Z]):/
        p.gsub('/', File::SEPARATOR)
      end

      # Normalize a file-system path into a stable cache key.  Different code
      # paths build the "same" file's path differently — e.g. a path built by
      # walking the on-disk directory tree during a full workspace scan versus
      # a path decoded from a `textDocument` URI sent by the editor, which on
      # Windows commonly differs only in the case of the drive letter (editors
      # conventionally lowercase it, `Dir.each_child` walks preserve on-disk
      # casing). Without normalizing, the same file gets indexed under two
      # different keys, producing duplicate call/type-hierarchy entries.
      private def normalize_key(path : String) : String
        key = path.gsub('\\', '/')
        if key =~ /^([a-zA-Z]):/
          key = key[0].downcase.to_s + key[1..]
        end
        key
      end
    end
  end
end
