defmodule Pyex.LexerTest do
  use ExUnit.Case, async: true

  alias Pyex.Lexer

  describe "atoms and operators" do
    test "integer literals" do
      assert {:ok, [{:integer, 1, 42}]} = Lexer.tokenize("42")
    end

    test "float literals" do
      assert {:ok, [{:float, 1, 3.14}]} = Lexer.tokenize("3.14")
    end

    test "identifiers" do
      assert {:ok, [{:name, 1, "foo"}]} = Lexer.tokenize("foo")
    end

    test "keywords are rewritten from names" do
      assert {:ok, [{:keyword, 1, "def"}]} = Lexer.tokenize("def")
      assert {:ok, [{:keyword, 1, "return"}]} = Lexer.tokenize("return")
      assert {:ok, [{:keyword, 1, "True"}]} = Lexer.tokenize("True")
    end

    test "multi-character operators come before single-character" do
      assert {:ok, [{:op, 1, :double_star}]} = Lexer.tokenize("**")
      assert {:ok, [{:op, 1, :floor_div}]} = Lexer.tokenize("//")
      assert {:ok, [{:op, 1, :eq}]} = Lexer.tokenize("==")
      assert {:ok, [{:op, 1, :neq}]} = Lexer.tokenize("!=")
      assert {:ok, [{:op, 1, :lte}]} = Lexer.tokenize("<=")
      assert {:ok, [{:op, 1, :gte}]} = Lexer.tokenize(">=")
    end

    test "arithmetic expression tokens in order" do
      {:ok, tokens} = Lexer.tokenize("2 + 3 * 4")

      assert [
               {:integer, 1, 2},
               {:op, 1, :plus},
               {:integer, 1, 3},
               {:op, 1, :star},
               {:integer, 1, 4}
             ] = tokens
    end
  end

  describe "indentation" do
    test "simple indent and dedent" do
      {:ok, tokens} = Lexer.tokenize("if x:\n    y")

      assert [
               {:keyword, 1, "if"},
               {:name, 1, "x"},
               {:op, 1, :colon},
               :newline,
               :indent,
               {:name, 2, "y"},
               :dedent
             ] = tokens
    end

    test "function definition produces correct indent structure" do
      source = "def add(a, b):\n    return a + b\n\nadd(1, 2)"
      {:ok, tokens} = Lexer.tokenize(source)

      assert [
               {:keyword, 1, "def"},
               {:name, 1, "add"},
               {:op, 1, :lparen},
               {:name, 1, "a"},
               {:op, 1, :comma},
               {:name, 1, "b"},
               {:op, 1, :rparen},
               {:op, 1, :colon},
               :newline,
               :indent,
               {:keyword, 2, "return"},
               {:name, 2, "a"},
               {:op, 2, :plus},
               {:name, 2, "b"},
               :newline,
               :dedent,
               {:name, 4, "add"},
               {:op, 4, :lparen},
               {:integer, 4, 1},
               {:op, 4, :comma},
               {:integer, 4, 2},
               {:op, 4, :rparen}
             ] = tokens
    end

    test "nested indentation with double dedent" do
      source = "if a:\n    if b:\n        x\ny"
      {:ok, tokens} = Lexer.tokenize(source)

      assert [
               {:keyword, 1, "if"},
               {:name, 1, "a"},
               {:op, 1, :colon},
               :newline,
               :indent,
               {:keyword, 2, "if"},
               {:name, 2, "b"},
               {:op, 2, :colon},
               :newline,
               :indent,
               {:name, 3, "x"},
               :newline,
               :dedent,
               :dedent,
               {:name, 4, "y"}
             ] = tokens
    end
  end

  describe "line tracking" do
    test "tokens carry correct line numbers across newlines" do
      source = "x = 1\ny = 2\nz = 3"
      {:ok, tokens} = Lexer.tokenize(source)

      lines = for {_, line, _} <- tokens, do: line

      assert [1, 1, 1, 2, 2, 2, 3, 3, 3] = lines
    end
  end

  describe "string literals" do
    test "double-quoted string" do
      assert {:ok, [{:string, 1, "hello"}]} = Lexer.tokenize(~s("hello"))
    end

    test "single-quoted string" do
      assert {:ok, [{:string, 1, "world"}]} = Lexer.tokenize("'world'")
    end

    test "string with escape sequences" do
      assert {:ok, [{:string, 1, "line1\nline2"}]} = Lexer.tokenize(~s("line1\\nline2"))
    end

    test "empty string" do
      assert {:ok, [{:string, 1, ""}]} = Lexer.tokenize(~s(""))
    end
  end

  describe "brackets and braces" do
    test "square brackets" do
      {:ok, tokens} = Lexer.tokenize("[1, 2]")

      assert [
               {:op, 1, :lbracket},
               {:integer, 1, 1},
               {:op, 1, :comma},
               {:integer, 1, 2},
               {:op, 1, :rbracket}
             ] = tokens
    end

    test "curly braces" do
      {:ok, tokens} = Lexer.tokenize("{}")
      assert [{:op, 1, :lbrace}, {:op, 1, :rbrace}] = tokens
    end

    test "dot operator" do
      {:ok, tokens} = Lexer.tokenize("foo.bar")

      assert [
               {:name, 1, "foo"},
               {:op, 1, :dot},
               {:name, 1, "bar"}
             ] = tokens
    end

    test "ellipsis literal" do
      {:ok, tokens} = Lexer.tokenize("...")
      assert [{:op, 1, :ellipsis}] = tokens
    end
  end

  describe "import keyword" do
    test "import is recognized as keyword" do
      assert {:ok, [{:keyword, 1, "import"}, {:name, 1, "json"}]} = Lexer.tokenize("import json")
    end
  end

  describe "bracket newline suppression" do
    test "newlines inside parentheses are suppressed" do
      {:ok, tokens} =
        Lexer.tokenize("""
        foo(
          1,
          2
        )
        """)

      refute :newline in tokens
      refute :indent in tokens
      refute :dedent in tokens

      assert {:name, 1, "foo"} in tokens
      assert {:op, 1, :lparen} in tokens
      assert {:integer, 2, 1} in tokens
      assert {:integer, 3, 2} in tokens
      assert {:op, 4, :rparen} in tokens
    end

    test "newlines inside square brackets are suppressed" do
      {:ok, tokens} =
        Lexer.tokenize("""
        [
          1,
          2,
          3
        ]
        """)

      refute :newline in tokens
      refute :indent in tokens
    end

    test "newlines inside curly braces are suppressed" do
      {:ok, tokens} =
        Lexer.tokenize("""
        {
          "a": 1,
          "b": 2
        }
        """)

      refute :newline in tokens
      refute :indent in tokens
    end

    test "nested brackets suppress correctly" do
      {:ok, tokens} =
        Lexer.tokenize("""
        foo(
          [1,
           2],
          3
        )
        """)

      refute :newline in tokens
      refute :indent in tokens
    end

    test "newlines outside brackets still produce indent/dedent" do
      {:ok, tokens} =
        Lexer.tokenize("""
        if x:
          y = [
            1,
            2
          ]
        """)

      assert :newline in tokens
      assert :indent in tokens
    end
  end

  describe "comments" do
    test "full-line comment is ignored" do
      assert {:ok, [:newline, {:name, 2, "x"}]} = Lexer.tokenize("# this is a comment\nx")
    end

    test "inline comment is stripped" do
      assert {:ok, [{:name, 1, "x"}, {:op, 1, :assign}, {:integer, 1, 1}]} =
               Lexer.tokenize("x = 1 # set x")
    end

    test "hash inside double-quoted string is preserved" do
      assert {:ok, [{:string, 1, "a#b"}]} = Lexer.tokenize(~s("a#b"))
    end

    test "hash inside single-quoted string is preserved" do
      assert {:ok, [{:string, 1, "a#b"}]} = Lexer.tokenize("'a#b'")
    end

    test "comment after string with hash inside" do
      {:ok, tokens} = Lexer.tokenize(~s(x = "a#b" # comment))

      assert [
               {:name, 1, "x"},
               {:op, 1, :assign},
               {:string, 1, "a#b"}
             ] = tokens
    end

    test "multiple comment lines" do
      source = "# first\n# second\nx = 1"
      {:ok, tokens} = Lexer.tokenize(source)

      names = for {:name, _, val} <- tokens, do: val
      assert "x" in names
    end

    test "comment-only source produces empty token list" do
      assert {:ok, []} = Lexer.tokenize("# just a comment")
    end

    test "comment preserves indentation structure" do
      source = "if x:\n    # a comment\n    y = 1"
      {:ok, tokens} = Lexer.tokenize(source)
      assert :indent in tokens
      assert {:name, _, "y"} = Enum.find(tokens, fn t -> match?({:name, _, "y"}, t) end)
    end
  end

  describe "error cases" do
    test "returns error for unrecognized characters" do
      assert {:error, "Lexer error:" <> _} = Lexer.tokenize("`")
    end
  end

  describe "prefixed numeric literals" do
    test "hex literals" do
      assert {:ok, [{:integer, 1, 255}]} = Lexer.tokenize("0xFF")
      assert {:ok, [{:integer, 1, 255}]} = Lexer.tokenize("0xff")
      assert {:ok, [{:integer, 1, 0}]} = Lexer.tokenize("0x0")
      assert {:ok, [{:integer, 1, 26}]} = Lexer.tokenize("0x1a")
    end

    test "octal literals" do
      assert {:ok, [{:integer, 1, 8}]} = Lexer.tokenize("0o10")
      assert {:ok, [{:integer, 1, 0}]} = Lexer.tokenize("0o0")
      assert {:ok, [{:integer, 1, 63}]} = Lexer.tokenize("0o77")
    end

    test "binary literals" do
      assert {:ok, [{:integer, 1, 10}]} = Lexer.tokenize("0b1010")
      assert {:ok, [{:integer, 1, 0}]} = Lexer.tokenize("0b0")
      assert {:ok, [{:integer, 1, 7}]} = Lexer.tokenize("0b111")
    end

    test "underscore separators in integers" do
      assert {:ok, [{:integer, 1, 1_000_000}]} = Lexer.tokenize("1_000_000")
      assert {:ok, [{:integer, 1, 42}]} = Lexer.tokenize("4_2")
    end

    test "underscore separators in floats" do
      assert {:ok, [{:float, 1, 1_000.5}]} = Lexer.tokenize("1_000.5")
    end

    test "underscore separators in hex" do
      assert {:ok, [{:integer, 1, 0xFFFF}]} = Lexer.tokenize("0xFF_FF")
    end

    test "underscore separators in binary" do
      assert {:ok, [{:integer, 1, 0b11110000}]} = Lexer.tokenize("0b1111_0000")
    end
  end

  describe "line continuation" do
    test "backslash joins lines" do
      {:ok, tokens} = Lexer.tokenize("x = 1 +\\\n2")
      names = for {:name, _, n} <- tokens, do: n
      assert "x" in names
    end

    test "backslash in expression" do
      assert Pyex.run!("x = 1 +\\\n2\nx") == 3
    end

    test "backslash in assignment" do
      result =
        Pyex.run!("""
        long_name =\
          42
        long_name
        """)

      assert result == 42
    end
  end

  describe "raw strings" do
    test "raw double-quoted string preserves backslashes" do
      assert {:ok, [{:string, 1, "\\n\\t"}]} = Lexer.tokenize(~S|r"\n\t"|)
    end

    test "raw single-quoted string preserves backslashes" do
      assert {:ok, [{:string, 1, "\\n\\t"}]} = Lexer.tokenize(~S|r'\n\t'|)
    end

    test "raw string with regex pattern" do
      assert {:ok, [{:string, 1, "\\d+\\.\\d+"}]} = Lexer.tokenize(~S|r"\d+\.\d+"|)
    end

    test "raw string with backslash-quote" do
      assert {:ok, [{:string, 1, "it\\'s"}]} = Lexer.tokenize(~S|r"it\'s"|)
    end

    test "raw string end-to-end" do
      assert Pyex.run!(~S|len(r"\n")|) == 2
    end

    test "raw string used for regex" do
      result =
        Pyex.run!("""
        import re
        pattern = r"\\d+"
        m = re.findall(pattern, "abc123def456")
        m
        """)

      assert result == ["123", "456"]
    end
  end

  describe "raw f-strings" do
    test "rf prefix emits an fstring token with the expression intact" do
      assert {:ok, [{:fstring, 1, "val {x}"}]} = Lexer.tokenize(~S|rf"val {x}"|)
    end

    test "fr prefix (reversed order) emits an fstring token" do
      assert {:ok, [{:fstring, 1, "val {x}"}]} = Lexer.tokenize(~S|fr"val {x}"|)
    end

    test "single-quoted rf prefix emits an fstring token" do
      assert {:ok, [{:fstring, 1, "val {x}"}]} = Lexer.tokenize(~S|rf'val {x}'|)
    end

    test "every case ordering of the combined prefix lexes identically" do
      for prefix <- ~w(rf fr Rf rF Fr fR RF FR) do
        assert {:ok, [{:fstring, 1, "val {x}"}]} =
                 Lexer.tokenize(prefix <> ~S|"val {x}"|),
               "prefix #{prefix} did not lex as a raw f-string"
      end
    end

    test "raw escape semantics: backslashes are preserved, unlike a plain f-string" do
      assert {:ok, [{:fstring, 1, "\\n{x}"}]} = Lexer.tokenize(~S|rf"\n{x}"|)
      assert {:ok, [{:fstring, 1, "\n{x}"}]} = Lexer.tokenize(~S|f"\n{x}"|)
    end

    test "backslash-quote does not terminate the raw f-string" do
      assert {:ok, [{:fstring, 1, "it\\'s {x}"}]} = Lexer.tokenize(~S|rf"it\'s {x}"|)
    end

    test "interpolates an expression end-to-end" do
      assert Pyex.run!("x = 5\nrf'val {x}'") == "val 5"
    end

    test "keeps backslashes literal while interpolating end-to-end" do
      assert Pyex.run!(~S|x = 5| <> "\n" <> ~S|rf"\d+ {x}"|) == "\\d+ 5"
    end
  end

  describe "string-prefix case-insensitivity" do
    test "uppercase F is an f-string, not an identifier plus string" do
      assert {:ok, [{:fstring, 1, "val {x}"}]} = Lexer.tokenize(~S|F"val {x}"|)
    end

    test "uppercase R is a raw string, not an identifier plus string" do
      assert {:ok, [{:string, 1, "\\n"}]} = Lexer.tokenize(~S|R"\n"|)
    end

    test "uppercase F interpolates end-to-end" do
      assert Pyex.run!("x = 5\nF'val {x}'") == "val 5"
    end
  end

  describe "unicode (u/U) prefix" do
    test "lowercase u prefix lexes as a plain string" do
      assert {:ok, [{:string, 1, "café"}]} = Lexer.tokenize(~S|u"café"|)
    end

    test "uppercase U prefix lexes as a plain string" do
      assert {:ok, [{:string, 1, "ab"}]} = Lexer.tokenize(~S|U'ab'|)
    end

    test "u prefix processes escapes like a normal string (it is a no-op)" do
      assert Pyex.run!(~S|u"a\nb"|) == "a\nb"
    end
  end

  describe "raw triple-quoted strings" do
    test "raw triple-quoted string preserves backslashes" do
      assert {:ok, [{:string, 1, "\\d+\\.\\d+"}]} = Lexer.tokenize(~S|r"""\d+\.\d+"""|)
    end

    test "raw triple-quoted string keeps embedded lone quotes" do
      assert {:ok, [{:string, 1, ~S|he said "hi" ok|}]} =
               Lexer.tokenize(~S|r"""he said "hi" ok"""|)
    end

    test "raw-f triple-quoted string keeps backslashes and interpolates" do
      assert Pyex.run!("x = 5\n" <> ~S|rf"""\d+ {x}"""|) == "\\d+ 5"
    end

    test "uppercase F triple-quoted string interpolates" do
      assert Pyex.run!("x = 5\n" <> ~S|F"""val {x}"""|) == "val 5"
    end
  end

  describe "byte-string prefixes" do
    test "plain bytes literal lexes to a bytes token" do
      assert {:ok, [{:bytes_literal, 1, "ab"}]} = Lexer.tokenize(~S|b"ab"|)
    end

    test "uppercase B and reversed/raw combos all lex to bytes tokens" do
      for prefix <- ~w(b B rb br Rb rB bR Br RB BR) do
        assert {:ok, [{:bytes_literal, 1, "ab"}]} = Lexer.tokenize(prefix <> ~S|"ab"|),
               "prefix #{prefix} did not lex as a bytes literal"
      end
    end

    test "cooked bytes interpret escapes; raw bytes keep them literal" do
      assert {:ok, [{:bytes_literal, 1, "a\nb"}]} = Lexer.tokenize(~S|b"a\nb"|)
      assert {:ok, [{:bytes_literal, 1, "a\\nb"}]} = Lexer.tokenize(~S|rb"a\nb"|)
    end

    test "bytes \\xhh is a single raw byte, not a UTF-8 codepoint" do
      assert {:ok, [{:bytes_literal, 1, <<0xFF>>}]} = Lexer.tokenize(~S|b"\xff"|)
    end

    test "bytes do not interpret \\u escapes (kept literal)" do
      # Source is b"é"; the \u stays literal in a bytes literal.
      assert {:ok, [{:bytes_literal, 1, "\\u00e9"}]} = Lexer.tokenize("b\"\\u00e9\"")
    end

    test "raw bytes triple-quoted keeps backslashes and embedded quotes" do
      assert {:ok, [{:bytes_literal, 1, ~S|a "b" \d|}]} =
               Lexer.tokenize(~S|rb"""a "b" \d"""|)
    end

    test "repr switches to double quotes when bytes contain a single quote" do
      assert Pyex.run!(~S|repr(b"say 'hi'")|) == ~S|b"say 'hi'"|
    end

    test "repr keeps single quotes and escapes when bytes contain a double quote" do
      assert Pyex.run!(~S|repr(b'say "hi"')|) == ~S|b'say "hi"'|
    end
  end

  describe "escape sequences" do
    test "carriage return \\r" do
      assert {:ok, [{:string, 1, "\r"}]} = Lexer.tokenize(~S|"\r"|)
    end

    test "null \\0" do
      assert {:ok, [{:string, 1, <<0>>}]} = Lexer.tokenize(~S|"\0"|)
    end

    test "bell \\a" do
      assert {:ok, [{:string, 1, "\a"}]} = Lexer.tokenize(~S|"\a"|)
    end

    test "backspace \\b" do
      assert {:ok, [{:string, 1, "\b"}]} = Lexer.tokenize(~S|"\b"|)
    end

    test "form feed \\f" do
      assert {:ok, [{:string, 1, "\f"}]} = Lexer.tokenize(~S|"\f"|)
    end

    test "vertical tab \\v" do
      assert {:ok, [{:string, 1, "\v"}]} = Lexer.tokenize(~S|"\v"|)
    end

    test "hex escape \\xNN" do
      assert {:ok, [{:string, 1, "A"}]} = Lexer.tokenize(~S|"\x41"|)
      {:ok, [{:string, 1, result}]} = Lexer.tokenize(~S|"\xff"|)
      assert result == <<0xFF::utf8>>
    end

    test "unicode escape \\uNNNN" do
      assert {:ok, [{:string, 1, "€"}]} = Lexer.tokenize(~S|"\u20AC"|)
      assert {:ok, [{:string, 1, "A"}]} = Lexer.tokenize(~S|"\u0041"|)
    end

    test "unicode escape \\UNNNNNNNN" do
      assert {:ok, [{:string, 1, "😀"}]} = Lexer.tokenize(~S|"\U0001F600"|)
    end

    test "unknown escape passes through" do
      assert {:ok, [{:string, 1, "\\d"}]} = Lexer.tokenize(~S|"\d"|)
      assert {:ok, [{:string, 1, "\\w"}]} = Lexer.tokenize(~S|"\w"|)
    end

    test "escape sequences in single-quoted strings" do
      assert {:ok, [{:string, 1, "\r"}]} = Lexer.tokenize(~S|'\r'|)
      assert {:ok, [{:string, 1, "A"}]} = Lexer.tokenize(~S|'\x41'|)
    end

    test "escape sequences in triple-quoted strings" do
      assert {:ok, [{:string, 1, "\r"}]} = Lexer.tokenize(~s|""\"\\r\"""|)
    end

    test "mixed escapes in one string" do
      assert {:ok, [{:string, 1, "a\tb\nc"}]} = Lexer.tokenize(~S|"a\tb\nc"|)
    end

    test "escape sequences work end-to-end" do
      assert Pyex.run!(~S|len("\r\n")|) == 2
      assert Pyex.run!(~S|"\x48\x65\x6c\x6c\x6f"|) == "Hello"
    end
  end
end
