defmodule Pyex do
  @moduledoc """
  A Python 3 interpreter written in Elixir.

  Pyex lexes, parses, and evaluates Python source code entirely
  within the BEAM -- no external runtime, no ports, and no Python
  runtime dependency.
  It is designed as a capabilities-based sandbox for running
  LLM-generated compute safely: every I/O operation (network,
  filesystem, database) is denied by default and must be
  explicitly granted through `Pyex.Ctx` options.

  ## Quick start

      Pyex.run!("sorted([3, 1, 2])")
      # => [1, 2, 3]

      {:ok, 42, _ctx} = Pyex.run("40 + 2")

  ## Public API

  - `compile/1` -- parse source to AST (reusable)
  - `run/2` -- execute source or AST, returns `{:ok, value, ctx}` or `{:error, error}`
  - `run!/2` -- execute, returns value or raises
  - `output/1` -- extract print output from a context

  ## Sandbox

  All external access is configured through `Pyex.Ctx` options.
  Python code can only reach what you explicitly grant:

      Pyex.run(source,
        env: %{"API_KEY" => "sk-..."},
        timeout: 5_000,
        modules: %{"mylib" => %{"greet" => {:builtin, fn [n] -> "hi \#{n}" end}}})

  See `run/2` for the full list of options.
  """

  alias Pyex.{Builtins, Ctx, Error, Lexer, Parser, Interpreter}

  @doc """
  Compiles a Python source string to an AST.

  The lexer and parser are pure functions of the source text.
  Callers that execute the same program repeatedly should
  compile once and reuse the AST with `run/2`.
  """
  @spec compile(String.t()) :: {:ok, Parser.ast_node()} | {:error, String.t()}
  def compile(source) when is_binary(source) do
    with {:ok, tokens} <- Lexer.tokenize(source),
         {:ok, ast} <- Parser.parse(tokens) do
      {:ok, Parser.ScopeResolve.resolve(ast)}
    end
  end

  @doc """
  Runs Python code through the full pipeline.

  Accepts either a source string or a pre-compiled AST.
  The optional second argument can be a `Pyex.Ctx` struct
  or a keyword list of options (forwarded to `Pyex.Ctx.new/1`).

  Returns `{:ok, value, ctx}` on success, or `{:error, reason}`.

  ## Options (when passing keyword list)

  - `:modules` -- custom Python modules available via `import`
  - `:filesystem` -- a filesystem backend struct (module derived automatically)
  - `:env` -- environment variables for `os.environ`
  - `:timeout` -- compute time budget in milliseconds
  - `:limits` -- resource limits as a keyword list or `Pyex.Limits` struct.
    Supported keys: `:timeout`, `:max_steps`, `:max_memory_bytes`,
    `:max_output_bytes`. When `limits: [timeout: N]` is provided,
    it supersedes a top-level `:timeout`. Ceilings are **on by
    default** (10M steps / 50 MB / 1 MB output); unspecified keys take
    their safe defaults, so partial limits are additive. Pass
    `limits: :none` to lift every ceiling for trusted code. See
    `Pyex.Limits` for details.
  - `:network` -- network access policy for the `requests` module.
    A list of rule maps, each with `:allowed_url_prefix` or
    `:dangerously_allow_full_internet_access`, plus optional `:methods`
    (default `["GET", "HEAD"]`) and `:headers` (injected into matching
    requests). When omitted, all network access is denied.
  - `:capabilities` -- list of enabled I/O capabilities (e.g.
    `[:boto3, :sql]`). All capabilities are denied by default.
  - `:boto3` -- shorthand for adding `:boto3` to capabilities.
  - `:sql` -- shorthand for adding `:sql` to capabilities.

  ## Return values

  Python values are mapped to Elixir at the boundary. Containers
  whose semantics match an Elixir native type are unwrapped; the
  rest are returned in a tagged form so the Python type is not
  lost.

  | Python type | Elixir shape |
  | ----------- | ------------ |
  | `int`, `float`, `str`, `bool`, `None` | scalar (`integer`, `float`, `binary`, `boolean`, `nil`) |
  | `list` | `[v1, v2, ...]` |
  | `dict` | `%{k1 => v1, ...}` |
  | `tuple` | `{:tuple, [v1, v2, ...]}` |
  | `set` | `{:set, MapSet.t()}` |
  | `frozenset` | `{:frozenset, MapSet.t()}` |
  | `range` | `{:range, start, stop, step}` |
  | `generator` | `{:generator, [v1, ...]}` (materialised) |
  | class instance | `{:instance, class, fields}` |

  Tagged shapes are used where Elixir's native type would lose
  information — for example, a Python tuple is *distinct* from
  a Python list, so collapsing both into `[...]` would be lossy.
  Container values are converted recursively, so a list of
  Python tuples surfaces as `[{:tuple, [...]}, {:tuple, [...]}]`.

  ## Examples

      {:ok, 42, _ctx} = Pyex.run("40 + 2")

      {:ok, ast} = Pyex.compile("40 + 2")
      {:ok, 42, _ctx} = Pyex.run(ast)

      {:ok, "hello world", _ctx} = Pyex.run(
        ~s|import mylib; mylib.greet("world")|,
        modules: %{
          "mylib" => %{
            "greet" => {:builtin, fn [name] -> "hello " <> name end}
          }
        })

      # Python tuple stays tagged so callers can pattern-match on it:
      {:ok, {:tuple, [1, 2, 3]}, _ctx} = Pyex.run("(1, 2, 3)")

      # List of tuples — each tuple is tagged:
      {:ok, [{:tuple, [1, 2]}, {:tuple, [3, 4]}], _ctx} =
        Pyex.run("[(1, 2), (3, 4)]")
  """
  @spec run(String.t() | Parser.ast_node(), Ctx.t() | keyword()) ::
          {:ok, Interpreter.pyvalue(), Ctx.t()}
          | {:error, Error.t()}
  def run(source_or_ast, ctx_or_opts \\ [])

  def run(source, ctx_or_opts) when is_binary(source) do
    case compile(source) do
      {:ok, ast} -> run(ast, ctx_or_opts)
      {:error, msg} -> {:error, Error.syntax(msg)}
    end
  end

  def run(ast, %Ctx{} = ctx) when is_tuple(ast) do
    start_mono = System.monotonic_time()

    :telemetry.execute([:pyex, :run, :start], %{system_time: System.system_time()}, %{})

    ctx = %Ctx{ctx | compute: 0.0, compute_started_at: System.monotonic_time()}

    # Match CPython's default decimal context (precision 28, banker's rounding).
    # The Elixir Decimal library defaults to :half_up; CPython uses :half_even.
    Decimal.Context.set(%Decimal.Context{precision: 28, rounding: :half_even})

    env = Builtins.runtime_env(ctx)

    result = Interpreter.run_with_ctx_result(ast, env, ctx)

    case result do
      {:ok, value, _env, final_ctx} ->
        final_ctx = close_open_handles(final_ctx)

        duration_ms =
          System.convert_time_unit(System.monotonic_time() - start_mono, :native, :microsecond) /
            1000.0

        :telemetry.execute([:pyex, :run, :stop], %{duration_ms: duration_ms}, %{
          compute: Ctx.compute_time(final_ctx)
        })

        derefed = Ctx.deep_deref(final_ctx, value)

        {:ok, Interpreter.Helpers.to_python_view(derefed),
         %{final_ctx | duration_ms: duration_ms}}

      {:error, msg, final_ctx} ->
        close_open_handles(final_ctx)

        duration_ms =
          System.convert_time_unit(System.monotonic_time() - start_mono, :native, :microsecond) /
            1000.0

        error = Error.from_message(msg)

        :telemetry.execute([:pyex, :run, :exception], %{duration_ms: duration_ms}, %{
          error: error
        })

        {:error, error}
    end
  end

  def run(ast, opts) when is_tuple(ast) and is_list(opts) do
    run(ast, Ctx.new(opts))
  end

  @doc """
  Runs Python code and returns the result directly.

  Raises on lexer, parser, or runtime errors. Accepts the
  same arguments as `run/2`.

  ## Examples

      42 = Pyex.run!("40 + 2")
      "hello" = Pyex.run!(ast, ctx)
  """
  @spec run!(String.t() | Parser.ast_node(), Ctx.t() | keyword()) :: Interpreter.pyvalue()
  def run!(source_or_ast, ctx_or_opts \\ []) do
    case run(source_or_ast, ctx_or_opts) do
      {:ok, result, _ctx} -> result
      {:error, %Error{message: msg}} -> raise msg
    end
  end

  @doc """
  Returns all captured print output as a single string.

  Python's `print()` records each line as an `:output` event
  in the context. This function extracts and joins them with newlines.

  ## Example

      {:ok, _val, ctx} = Pyex.run("print('hello')")
      "hello" = Pyex.output(ctx)

      {:ok, _val, ctx} = Pyex.run("print('line1')\nprint('line2')")
      "line1\nline2" = Pyex.output(ctx)
  """
  @spec output(Ctx.t()) :: String.t()
  def output(%Ctx{} = ctx), do: ctx |> Ctx.output() |> IO.iodata_to_binary()

  @spec close_open_handles(Ctx.t()) :: Ctx.t()
  defp close_open_handles(final_ctx) do
    Enum.reduce(Map.keys(final_ctx.handles), final_ctx, fn id, ctx ->
      case Ctx.close_handle(ctx, id) do
        {:ok, ctx} -> ctx
        {:error, _} -> ctx
      end
    end)
  end
end
