defmodule Pyex.Interpreter.Import do
  @moduledoc """
  Module resolution and import logic for the Pyex interpreter.

  Resolves Python `import` and `from ... import` statements against
  the stdlib registry, custom modules, and the pluggable filesystem.
  Supports dotted module names (e.g. `os.path`) and caches imported
  filesystem modules in `ctx.imported_modules`.

  Modules are represented as `{:module, name, attrs_map}` pyvalues so
  that `type(m)` returns `"module"` and `repr(m)` renders as
  `<module 'name'>`.  The attrs map may contain nested `{:module, ...}`
  values to expose submodules (e.g. `os.path`).
  """

  alias Pyex.{Builtins, Ctx, Env, Interpreter}

  @doc """
  Evaluates an `import ...` statement.
  """
  @spec eval_import([String.t() | {String.t(), String.t() | nil}], Env.t(), Ctx.t()) ::
          {nil, Env.t(), Ctx.t()} | {{:exception, String.t()}, Env.t(), Ctx.t()}
  def eval_import(imports, env, ctx) when is_list(imports) do
    Enum.reduce_while(imports, {nil, env, ctx}, fn import_spec, {_, env, ctx} ->
      {module_name, alias_name} =
        case import_spec do
          {module_name, alias_name} -> {module_name, alias_name}
          module_name when is_binary(module_name) -> {module_name, nil}
        end

      bind_as = alias_name || import_binding_name(module_name)

      case resolve_module(module_name, env, ctx) do
        {:ok, module_value, ctx} ->
          bound_value = binding_value(module_name, alias_name, module_value, env, ctx)
          {:cont, {nil, Env.put(env, bind_as, bound_value), ctx}}

        {:import_error, msg, ctx} ->
          {:halt, {{:exception, msg}, env, ctx}}

        {:unknown_module, ctx} ->
          {:halt,
           {{:exception,
             "ImportError: no module named '#{module_name}'#{import_hint(module_name)}"}, env,
            ctx}}
      end
    end)
  end

  @doc """
  Evaluates a `from ... import ...` statement.
  """
  @spec eval_from_import(String.t(), [{String.t(), String.t() | nil}], Env.t(), Ctx.t()) ::
          {nil, Env.t(), Ctx.t()} | {{:exception, String.t()}, Env.t(), Ctx.t()}
  # `from __future__ import annotations` (and the other historical feature
  # flags) are no-ops on Python 3 — their behavior is already the default.
  def eval_from_import("__future__", _names, env, ctx), do: {nil, env, ctx}

  def eval_from_import(module_name, names, env, ctx) do
    case resolve_module(module_name, env, ctx) do
      {:ok, module, ctx} ->
        attrs = module_attrs(module)

        Enum.reduce_while(names, {nil, env, ctx}, fn {name, alias_name}, {_, env, ctx} ->
          bind_as = alias_name || name

          case Map.fetch(attrs, name) do
            {:ok, value} ->
              {:cont, {nil, Env.put(env, bind_as, value), ctx}}

            :error ->
              {:halt,
               {{:exception, "ImportError: cannot import name '#{name}' from '#{module_name}'"},
                env, ctx}}
          end
        end)

      {:import_error, msg, ctx} ->
        {{:exception, msg}, env, ctx}

      {:unknown_module, ctx} ->
        {{:exception, "ImportError: no module named '#{module_name}'#{import_hint(module_name)}"},
         env, ctx}
    end
  end

  @doc """
  Returns the attribute map for a module value, regardless of whether
  it is a `{:module, ...}` pyvalue or a bare map (legacy shape).
  """
  @spec module_attrs(Interpreter.pyvalue() | map()) :: map()
  def module_attrs({:module, _name, attrs}), do: attrs
  def module_attrs(attrs) when is_map(attrs), do: attrs

  @doc """
  Returns a helpful error suffix for unknown module names.
  """
  @spec import_hint(String.t()) :: String.t()
  def import_hint(name) when name in ["urllib2", "http", "httplib", "httpx", "aiohttp"],
    do: ". Use 'import requests' instead"

  def import_hint("http." <> _), do: ". Use 'import requests' instead"

  def import_hint("sys"), do: ". Use 'import os' for environ access"
  def import_hint("pathlib"), do: ". Use the filesystem option for file operations"
  def import_hint("io"), do: ". Use the filesystem option for file operations"

  def import_hint(_) do
    names = ["os" | Pyex.Stdlib.module_names()] |> Enum.join(", ")
    ". Available modules: #{names}"
  end

  @spec import_binding_name(String.t()) :: String.t()
  defp import_binding_name(module_name) do
    module_name
    |> String.split(".")
    |> hd()
  end

  @spec binding_value(
          String.t(),
          String.t() | nil,
          Interpreter.pyvalue(),
          Env.t(),
          Ctx.t()
        ) :: Interpreter.pyvalue()
  defp binding_value(module_name, alias_name, module_value, env, ctx) do
    cond do
      alias_name != nil ->
        module_value

      String.contains?(module_name, ".") ->
        root = import_binding_name(module_name)
        {:ok, root_module, _ctx} = resolve_single_module(root, env, ctx)
        root_module

      true ->
        module_value
    end
  end

  @doc """
  Resolves a module name to its `{:module, name, attrs}` pyvalue.

  Returns `{:ok, module_pyvalue, ctx}`, `{:unknown_module, ctx}`,
  or `{:import_error, message, ctx}`.
  """
  @spec resolve_module(String.t(), Env.t(), Ctx.t()) ::
          {:ok, Interpreter.pyvalue(), Ctx.t()}
          | {:unknown_module, Ctx.t()}
          | {:import_error, String.t(), Ctx.t()}
  def resolve_module(name, env, ctx) do
    case String.split(name, ".") do
      [single] ->
        resolve_single_module(single, env, ctx)

      [root | parts] ->
        # Dotted names have two possible meanings:
        #   1. A nested submodule walk on a module (e.g. `os.path`) — this
        #      is how the stdlib exposes submodules.
        #   2. A file in a package directory on the virtual filesystem
        #      (e.g. `pkg.math_utils` → `pkg/math_utils.py`).
        # Try the submodule walk first to preserve stdlib precedence,
        # then fall back to a direct filesystem lookup so user code can
        # organize modules into package directories.
        case resolve_single_module(root, env, ctx) do
          {:ok, module_value, ctx} ->
            case resolve_dotted_parts(module_value, parts, name, ctx) do
              {:ok, _, _} = ok -> ok
              {:unknown_module, ctx} -> resolve_filesystem_module(name, env, ctx)
            end

          {:unknown_module, ctx} ->
            resolve_filesystem_module(name, env, ctx)

          other ->
            other
        end
    end
  end

  @doc """
  Re-imports `name`, bypassing the per-run import cache — the substrate for
  `importlib.reload`. Evicts `ctx.imported_modules[name]` so a filesystem
  module is re-read and re-executed from its (possibly edited) source, then
  resolves it afresh. Same return shape as `resolve_module/3`.
  """
  @spec reload_module(String.t(), Env.t(), Ctx.t()) ::
          {:ok, Interpreter.pyvalue(), Ctx.t()}
          | {:unknown_module, Ctx.t()}
          | {:import_error, String.t(), Ctx.t()}
  def reload_module(name, env, ctx) do
    ctx = %{ctx | imported_modules: Map.delete(ctx.imported_modules, name)}
    resolve_module(name, env, ctx)
  end

  @spec resolve_single_module(String.t(), Env.t(), Ctx.t()) ::
          {:ok, Interpreter.pyvalue(), Ctx.t()}
          | {:unknown_module, Ctx.t()}
          | {:import_error, String.t(), Ctx.t()}
  defp resolve_single_module(name, env, ctx) do
    case Map.fetch(ctx.modules, name) do
      {:ok, value} ->
        {:ok, to_module_value(name, value), ctx}

      :error ->
        case resolve_builtin_module(name, ctx) do
          {:ok, value} ->
            {:ok, to_module_value(name, value), ctx}

          :unknown_module ->
            resolve_filesystem_module(name, env, ctx)
        end
    end
  end

  @spec to_module_value(String.t(), Interpreter.pyvalue() | map()) :: Interpreter.pyvalue()
  defp to_module_value(_name, {:module, _, _} = mod), do: mod
  defp to_module_value(name, attrs) when is_map(attrs), do: {:module, name, attrs}

  @spec resolve_dotted_parts(
          {:module, String.t(), %{optional(String.t()) => Interpreter.pyvalue()}},
          [String.t()],
          String.t(),
          Ctx.t()
        ) ::
          {:ok, Interpreter.pyvalue(), Ctx.t()} | {:unknown_module, Ctx.t()}
  defp resolve_dotted_parts(value, [], _full_name, ctx), do: {:ok, value, ctx}

  defp resolve_dotted_parts({:module, _, attrs}, [part | rest], full_name, ctx) do
    case Map.fetch(attrs, part) do
      {:ok, {:module, _, _} = sub} ->
        resolve_dotted_parts(sub, rest, full_name, ctx)

      {:ok, _} ->
        {:unknown_module, ctx}

      :error ->
        {:unknown_module, ctx}
    end
  end

  @spec resolve_builtin_module(String.t(), Ctx.t()) ::
          {:ok, Interpreter.pyvalue()} | :unknown_module
  defp resolve_builtin_module("os", ctx) do
    path_attrs = %{
      "join" => {:builtin, &os_path_join/1},
      "exists" => {:builtin, &os_path_exists/1},
      "basename" => {:builtin, &os_path_basename/1},
      "dirname" => {:builtin, &os_path_dirname/1},
      "splitext" => {:builtin, &os_path_splitext/1},
      "isfile" => {:builtin, &os_path_isfile/1},
      "isdir" => {:builtin, &os_path_isdir/1}
    }

    attrs = %{
      "environ" => ctx.env,
      "path" => {:module, "os.path", path_attrs},
      "getcwd" => {:builtin, &os_getcwd/1},
      "chdir" => {:builtin, &os_chdir/1},
      "makedirs" => {:builtin, &os_makedirs/1},
      "listdir" => {:builtin, &os_listdir/1},
      "walk" => {:builtin, &os_walk/1},
      "remove" => {:builtin, &os_remove/1},
      "unlink" => {:builtin, &os_remove/1}
    }

    {:ok, {:module, "os", attrs}}
  end

  defp resolve_builtin_module(name, _ctx) do
    case Pyex.Stdlib.fetch(name) do
      {:ok, attrs} -> {:ok, {:module, name, attrs}}
      :unknown_module -> :unknown_module
    end
  end

  @spec os_listdir([Interpreter.pyvalue()]) ::
          {:ctx_call, (Pyex.Env.t(), Ctx.t() -> {term(), Pyex.Env.t(), Ctx.t()})}
          | {:exception, String.t()}
  defp os_listdir([]) do
    os_listdir([""])
  end

  defp os_listdir([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {:ctx_call,
         fn env, ctx ->
           case ctx.filesystem do
             nil ->
               {{:exception, "OSError: no filesystem configured"}, env, ctx}

             fs ->
               case Pyex.Path.list_dir(fs, ctx.cwd, path) do
                 {:ok, entries, fs} ->
                   {entries, env, %{ctx | filesystem: fs}}

                 {:error, msg} ->
                   {{:exception, msg}, env, ctx}
               end
           end
         end}

      :error ->
        {:exception, "TypeError: listdir: path should be string"}
    end
  end

  defp os_listdir(_args) do
    {:exception, "TypeError: listdir expected at most 1 argument"}
  end

  @spec os_remove([Interpreter.pyvalue()]) ::
          {:ctx_call, (Pyex.Env.t(), Ctx.t() -> {term(), Pyex.Env.t(), Ctx.t()})}
          | {:exception, String.t()}
  defp os_remove([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {:ctx_call,
         fn env, ctx ->
           case ctx.filesystem do
             nil ->
               {{:exception, "OSError: no filesystem configured"}, env, ctx}

             fs ->
               case Pyex.FS.rm(fs, ctx.cwd, path) do
                 {:ok, fs} -> {nil, env, %{ctx | filesystem: fs}}
                 {:error, msg} -> {{:exception, msg}, env, ctx}
               end
           end
         end}

      :error ->
        {:exception, "TypeError: remove: path should be string"}
    end
  end

  defp os_remove(_args) do
    {:exception, "TypeError: remove expected exactly 1 argument"}
  end

  @spec os_walk([Interpreter.pyvalue()]) ::
          {:ctx_call, (Pyex.Env.t(), Ctx.t() -> {term(), Pyex.Env.t(), Ctx.t()})}
          | {:exception, String.t()}
  defp os_walk([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {:ctx_call,
         fn env, ctx ->
           case ctx.filesystem do
             nil ->
               {{:exception, "OSError: no filesystem configured"}, env, ctx}

             fs ->
               case Pyex.Path.walk(fs, ctx.cwd, path) do
                 {:ok, entries, fs} ->
                   rows =
                     Enum.map(entries, fn {root, dirs, files} ->
                       {:tuple, [root, dirs, Enum.map(files, &Pyex.Path.basename/1)]}
                     end)

                   {token, ctx} = Ctx.new_iterator(%{ctx | filesystem: fs}, rows)
                   {token, env, ctx}

                 {:error, msg} ->
                   {{:exception, msg}, env, ctx}
               end
           end
         end}

      :error ->
        {:exception, "TypeError: walk: top should be string"}
    end
  end

  defp os_walk(_args), do: {:exception, "TypeError: walk expected exactly 1 argument"}

  @spec resolve_filesystem_module(String.t(), Env.t(), Ctx.t()) ::
          {:ok, Interpreter.pyvalue(), Ctx.t()}
          | {:unknown_module, Ctx.t()}
          | {:import_error, String.t(), Ctx.t()}
  defp resolve_filesystem_module(_name, _env, %{filesystem: nil} = ctx) do
    {:unknown_module, ctx}
  end

  defp resolve_filesystem_module(name, _env, ctx) do
    case Map.fetch(ctx.imported_modules, name) do
      {:ok, cached} ->
        {:ok, to_module_value(name, cached), ctx}

      :error ->
        # Module files resolve from the filesystem root (sys.path-style), not
        # the cwd, so an import means the same thing wherever the program is.
        path = String.replace(name, ".", "/") <> ".py"

        case Pyex.FS.read_file(ctx.filesystem, "/", path) do
          {:ok, source, fs} ->
            ctx = %{ctx | filesystem: fs}

            case Pyex.compile(source) do
              {:ok, ast} ->
                mod_env =
                  Builtins.env()
                  |> Env.push_scope()
                  |> Env.put("__name__", name)

                case Interpreter.eval(ast, mod_env, ctx) do
                  {{:exception, msg}, _mod_env, ctx} ->
                    {:import_error, "ImportError: error in '#{name}': #{msg}", ctx}

                  {_val, mod_env, ctx} ->
                    attrs = collect_module_bindings(mod_env)
                    module = {:module, name, attrs}

                    ctx = %{
                      ctx
                      | imported_modules: Map.put(ctx.imported_modules, name, module)
                    }

                    {:ok, module, ctx}
                end

              {:error, msg} ->
                {:import_error, "SyntaxError: error in '#{name}': #{msg}", ctx}
            end

          {:error, _} ->
            {:unknown_module, ctx}
        end
    end
  end

  @spec collect_module_bindings(Env.t()) :: %{optional(String.t()) => Interpreter.pyvalue()}
  defp collect_module_bindings(env) do
    env
    |> Env.all_bindings()
    |> Enum.reject(fn {name, _} -> String.starts_with?(name, "__") end)
    |> Map.new()
  end

  # os.path module functions
  defp os_path_join(args) when args != [] do
    case coerce_paths(args) do
      {:ok, parts} -> Pyex.Path.join(parts)
      {:error, msg} -> {:exception, msg}
    end
  end

  defp os_path_join(_args) do
    {:exception, "TypeError: join() takes at least 1 argument"}
  end

  defp os_path_exists([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {:ctx_call,
         fn env, ctx ->
           case ctx.filesystem do
             nil ->
               {false, env, ctx}

             fs ->
               {exists, fs} = Pyex.Path.exists?(fs, ctx.cwd, path)
               {exists, env, %{ctx | filesystem: fs}}
           end
         end}

      :error ->
        {:exception, "TypeError: path should be string"}
    end
  end

  defp os_path_basename([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} -> Pyex.Path.basename(path)
      :error -> {:exception, "TypeError: path should be string"}
    end
  end

  defp os_path_dirname([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} -> Pyex.Path.dirname(path)
      :error -> {:exception, "TypeError: path should be string"}
    end
  end

  defp os_path_splitext([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {root, ext} = Pyex.Path.splitext(path)
        {:tuple, [root, ext]}

      :error ->
        {:exception, "TypeError: path should be string"}
    end
  end

  defp os_path_isfile([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {:ctx_call,
         fn env, ctx ->
           case ctx.filesystem do
             nil ->
               {false, env, ctx}

             fs ->
               {is_file, fs} = Pyex.Path.file?(fs, ctx.cwd, path)
               {is_file, env, %{ctx | filesystem: fs}}
           end
         end}

      :error ->
        {:exception, "TypeError: path should be string"}
    end
  end

  defp os_path_isdir([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {:ctx_call,
         fn env, ctx ->
           case ctx.filesystem do
             nil ->
               {false, env, ctx}

             fs ->
               {is_dir, fs} = Pyex.Path.dir?(fs, ctx.cwd, path)
               {is_dir, env, %{ctx | filesystem: fs}}
           end
         end}

      :error ->
        {:exception, "TypeError: path should be string"}
    end
  end

  defp os_makedirs([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {:ctx_call,
         fn env, ctx ->
           case ctx.filesystem do
             nil ->
               {{:exception, "OSError: no filesystem configured"}, env, ctx}

             fs ->
               case Pyex.Path.mkdir_p(fs, ctx.cwd, path) do
                 {:ok, fs} -> {nil, env, %{ctx | filesystem: fs}}
                 {:error, msg} -> {{:exception, msg}, env, ctx}
               end
           end
         end}

      :error ->
        {:exception, "TypeError: path should be string"}
    end
  end

  defp os_makedirs([path, kwargs]) when is_map(kwargs) do
    case Map.keys(kwargs) -- ["exist_ok"] do
      [] ->
        os_makedirs([path])

      [name | _] ->
        {:exception, "TypeError: makedirs() got an unexpected keyword argument '#{name}'"}
    end
  end

  @spec os_getcwd([Interpreter.pyvalue()]) ::
          {:ctx_call, (Pyex.Env.t(), Ctx.t() -> {term(), Pyex.Env.t(), Ctx.t()})}
          | {:exception, String.t()}
  defp os_getcwd([]) do
    {:ctx_call, fn env, ctx -> {ctx.cwd, env, ctx} end}
  end

  defp os_getcwd(_args), do: {:exception, "TypeError: getcwd() takes no arguments"}

  @spec os_chdir([Interpreter.pyvalue()]) ::
          {:ctx_call, (Pyex.Env.t(), Ctx.t() -> {term(), Pyex.Env.t(), Ctx.t()})}
          | {:exception, String.t()}
  defp os_chdir([path]) do
    case Pyex.Path.coerce(path) do
      {:ok, path} ->
        {:ctx_call,
         fn env, ctx ->
           target = Pyex.FS.resolve(ctx.cwd, path)

           case ctx.filesystem do
             nil ->
               {nil, env, %{ctx | cwd: target}}

             fs ->
               case Pyex.FS.stat(fs, ctx.cwd, path) do
                 {:ok, :directory, fs} ->
                   {nil, env, %{ctx | cwd: target, filesystem: fs}}

                 {:ok, :regular, fs} ->
                   {{:exception, "NotADirectoryError: [Errno 20] Not a directory: '#{path}'"},
                    env, %{ctx | filesystem: fs}}

                 {:error, msg} ->
                   {{:exception, msg}, env, ctx}
               end
           end
         end}

      :error ->
        {:exception, "TypeError: chdir: path should be string"}
    end
  end

  defp os_chdir(_args), do: {:exception, "TypeError: chdir() takes exactly 1 argument"}

  @spec coerce_paths([Interpreter.pyvalue()]) :: {:ok, [String.t()]} | {:error, String.t()}
  defp coerce_paths(paths) do
    Enum.reduce_while(paths, {:ok, []}, fn path, {:ok, acc} ->
      case Pyex.Path.coerce(path) do
        {:ok, path} -> {:cont, {:ok, [path | acc]}}
        :error -> {:halt, {:error, "TypeError: path should be string"}}
      end
    end)
    |> case do
      {:ok, acc} -> {:ok, Enum.reverse(acc)}
      {:error, _} = error -> error
    end
  end
end
