package pythonlang

import (
	"strings"
	"testing"

	"rickub.com/turbo-editors/turbo-core/syntax"
)

// classAt returns the class covering a rune column on a line, and whether any
// span covers it at all. It is how nearly every test below asks its question.
func classAt(spans [][]syntax.Span, line, col int) (syntax.Class, bool) {
	if line < 0 || line >= len(spans) {
		return 0, false
	}
	for _, s := range spans[line] {
		if col >= s.Start && col < s.End {
			return s.Class, true
		}
	}
	return 0, false
}

// classOfFirst returns the class of the first occurrence of word in src.
func classOfFirst(t *testing.T, src, word string) syntax.Class {
	t.Helper()

	index := strings.Index(src, word)
	if index < 0 {
		t.Fatalf("%q does not appear in the source", word)
	}
	line := strings.Count(src[:index], "\n")
	col := index - (strings.LastIndex(src[:index], "\n") + 1)

	class, ok := classAt(Highlight(src), line, col)
	if !ok {
		t.Fatalf("no span covers %q at line %d column %d", word, line, col)
	}
	return class
}

// spanOfFirst returns the span covering the first occurrence of word, so that a
// test can check where a construct *ends* and not only what colour it is.
func spanOfFirst(t *testing.T, src, word string) syntax.Span {
	t.Helper()

	index := strings.Index(src, word)
	if index < 0 {
		t.Fatalf("%q does not appear in the source", word)
	}
	line := strings.Count(src[:index], "\n")
	col := index - (strings.LastIndex(src[:index], "\n") + 1)

	for _, s := range Highlight(src)[line] {
		if col >= s.Start && col < s.End {
			return s
		}
	}
	t.Fatalf("no span covers %q at line %d column %d", word, line, col)
	return syntax.Span{}
}

// --- the three invariants the editor relies on ------------------------------

func TestHighlightReturnsOneEntryPerLine(t *testing.T) {
	// The editor indexes the result by line number without checking, so a short
	// result is an index out of range in the middle of a redraw.
	tests := []struct {
		name string
		src  string
		want int
	}{
		{"empty", "", 1},
		{"one line without a terminator", "x = 1", 1},
		{"one line with a terminator", "x = 1\n", 2},
		{"three lines", "a\nb\nc", 3},
		{"an unterminated triple quote", `x = """a` + "\nb\nc", 3},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			if got := len(Highlight(tc.src)); got != tc.want {
				t.Errorf("Highlight(%q) returned %d lines, want %d", tc.src, got, tc.want)
			}
		})
	}
}

func TestSpansAreInOrderAndDoNotOverlap(t *testing.T) {
	// They are drawn in order, so two out of order paint over each other and
	// nothing fails. This has happened in this family, in a scanner that
	// emitted a quote after the name it belonged to.
	const src = `#!/usr/bin/env python3
"""A module docstring
spanning two lines."""

import re
from dataclasses import dataclass

MAX_SIZE = 1_000
PATTERN = re.compile(r"\d+(\.\d+)?")


@dataclass(frozen=True)
class Measurement:
    """One reading."""

    name: str
    value: float = 0.5

    def scaled(self, factor: float = 1.5e-3) -> float:
        return self.value * factor


def main() -> None:
    for line in open("input.txt"):
        match line.strip():
            case "":
                continue
            case other:
                print(f"{other!r} -> {len(other)}", end="")


if __name__ == "__main__":
    main()
`

	for line, spans := range Highlight(src) {
		previous := syntax.Span{End: -1}
		for _, span := range spans {
			switch {
			case span.Start < previous.End:
				t.Errorf("line %d: %v starts at %d, inside %v which ends at %d",
					line, span.Class, span.Start, previous.Class, previous.End)
			case span.Start >= span.End:
				t.Errorf("line %d: %v is empty at %d", line, span.Class, span.Start)
			}
			previous = span
		}
	}
}

func TestBrokenSourceStillColours(t *testing.T) {
	// Source under the cursor is invalid most of the time it is being typed. A
	// scanner that gives up is a scanner that flickers off.
	broken := []string{
		`x = "unterminated`,
		`x = '''unterminated`,
		"def f(",
		"class",
		"@",
		"f'{",
		"x = 0x",
		"    )))",
		"\\",
		"x = 1.2.3.4",
	}

	for _, src := range broken {
		t.Run(src, func(t *testing.T) {
			spans := Highlight(src)
			if len(spans) != 1 {
				t.Fatalf("Highlight(%q) returned %d lines, want 1", src, len(spans))
			}
			for _, span := range spans[0] {
				if span.Start < 0 || span.End > len([]rune(src)) {
					t.Errorf("%v runs from %d to %d, outside a line of %d runes",
						span.Class, span.Start, span.End, len([]rune(src)))
				}
			}
		})
	}
}

// --- one test per construct -------------------------------------------------

func TestEachTokenClass(t *testing.T) {
	const src = `# a comment
import os
from typing import Iterator


class Shape:
    def __init__(self, x: int = 0) -> None:
        self.x = x
        name = "world"
        ratio = 1.5
        count = 42
        missing = None
        ok = True
        print(len(name))
`

	tests := []struct {
		word string
		want syntax.Class
	}{
		{"# a comment", syntax.ClassComment},
		{"import", syntax.ClassKeyword},
		{"from", syntax.ClassKeyword},
		{"class", syntax.ClassKeyword},
		{"def", syntax.ClassKeyword},
		{"Shape", syntax.ClassType},
		{"int", syntax.ClassType},
		{"__init__", syntax.ClassBuiltin},
		{"self", syntax.ClassBuiltin},
		{`"world"`, syntax.ClassString},
		{"1.5", syntax.ClassNumber},
		{"42", syntax.ClassNumber},
		{"None", syntax.ClassConstant},
		{"True", syntax.ClassConstant},
		{"print", syntax.ClassBuiltin},
		{"len", syntax.ClassBuiltin},
		{":", syntax.ClassPunctuation},
		{"=", syntax.ClassOperator},
	}

	for _, test := range tests {
		t.Run(test.word, func(t *testing.T) {
			if got := classOfFirst(t, src, test.word); got != test.want {
				t.Errorf("%q is %v, want %v", test.word, got, test.want)
			}
		})
	}
}

func TestACallIsAFunction(t *testing.T) {
	if got := classOfFirst(t, "n = compute(3)", "compute"); got != syntax.ClassFunction {
		t.Errorf("compute is %v, want function", got)
	}
}

func TestADeclaredFunctionNameIsAFunction(t *testing.T) {
	if got := classOfFirst(t, "def parse(text):\n    pass\n", "parse"); got != syntax.ClassFunction {
		t.Errorf("parse is %v, want function", got)
	}
}

// A class is called exactly the way a function is, so the parenthesis cannot
// tell them apart and the naming convention has to.
func TestACapitalisedNameIsATypeEvenWhenItIsCalled(t *testing.T) {
	tests := []struct{ src, word string }{
		{`raise ValueError("nope")`, "ValueError"},
		{`thing = Measurement(1)`, "Measurement"},
		{`class Measurement:`, "Measurement"},
		{`def f() -> Measurement:`, "Measurement"},
	}

	for _, tc := range tests {
		t.Run(tc.src, func(t *testing.T) {
			if got := classOfFirst(t, tc.src, tc.word); got != syntax.ClassType {
				t.Errorf("%q in %q is %v, want type", tc.word, tc.src, got)
			}
		})
	}
}

func TestAWordInCapitalsIsAConstant(t *testing.T) {
	for _, word := range []string{"MAX_SIZE", "PI", "HTTP_PORT", "_PRIVATE", "V2"} {
		t.Run(word, func(t *testing.T) {
			src := word + " = 1"
			if got := classOfFirst(t, src, word); got != syntax.ClassConstant {
				t.Errorf("%s is %v, want constant", word, got)
			}
		})
	}
}

func TestASingleCapitalIsATypeNotAConstant(t *testing.T) {
	// The constant rule wants at least two runes, so that a one-letter type
	// variable — T, the name every generic in the standard library uses — is
	// not read as a constant.
	if got := classOfFirst(t, `T = TypeVar("T")`, "T"); got != syntax.ClassType {
		t.Errorf("T is %v, want type", got)
	}
}

func TestNumbersInEveryBaseAndShape(t *testing.T) {
	numbers := []string{
		"42", "1_000", "0xFF", "0o17", "0b1010", "1.5", ".5", "1.", "1e10",
		"1.5e-3", "1E+7", "3j", "0x_FF_FF",
	}

	for _, number := range numbers {
		t.Run(number, func(t *testing.T) {
			src := "x = " + number + "\n"
			span := spanOfFirst(t, src, number)
			if span.Class != syntax.ClassNumber {
				t.Errorf("%s is %v, want number", number, span.Class)
			}
			if got := span.End - span.Start; got != len([]rune(number)) {
				t.Errorf("%s is coloured over %d runes, want %d", number, got, len([]rune(number)))
			}
		})
	}
}

// 0xE-1 is a hexadecimal literal minus one. The E is a digit here, not the e of
// an exponent, so the sign after it is an operator and not part of the number.
func TestAHexadecimalLiteralDoesNotSwallowTheSignAfterIt(t *testing.T) {
	span := spanOfFirst(t, "x = 0xE-1", "0xE")

	if span.End != len("x = 0xE") {
		t.Errorf("0xE is coloured to column %d, want %d — the minus was taken as an exponent's sign",
			span.End, len("x = 0xE"))
	}
}

// A float has at most one dot. Without that limit, `1.2.3` is one long number
// and the version string somebody is halfway through typing paints the line.
func TestANumberStopsAtItsSecondDot(t *testing.T) {
	span := spanOfFirst(t, "x = 1.2.3", "1.2")

	if span.End != len("x = 1.2") {
		t.Errorf("1.2 is coloured to column %d, want %d — the second dot was swallowed",
			span.End, len("x = 1.2"))
	}
}

func TestADottedAttributeIsNotANumber(t *testing.T) {
	if got := classOfFirst(t, "value = thing.count", "count"); got != syntax.ClassIdentifier {
		t.Errorf("count after a dot is %v, want identifier", got)
	}
}

// --- strings ----------------------------------------------------------------

func TestEveryStringPrefixOpensAString(t *testing.T) {
	prefixes := []string{"", "r", "R", "b", "B", "u", "U", "f", "F", "rb", "br", "fr", "rf", "Rb", "BR"}

	for _, prefix := range prefixes {
		t.Run("prefix "+prefix, func(t *testing.T) {
			literal := prefix + `"hello"`
			src := "x = " + literal
			span := spanOfFirst(t, src, literal)
			if span.Class != syntax.ClassString {
				t.Errorf("%s is %v, want string", literal, span.Class)
			}
			if got := span.End - span.Start; got != len([]rune(literal)) {
				t.Errorf("%s is coloured over %d runes, want %d — the prefix was left out",
					literal, got, len([]rune(literal)))
			}
		})
	}
}

func TestAnIdentifierEndingInAPrefixLetterIsNotAString(t *testing.T) {
	// foo"bar" must be an identifier and a string, not one run: the prefix
	// letters are only a prefix when nothing but them comes before the quote.
	if got := classOfFirst(t, `foo"bar"`, "foo"); got != syntax.ClassIdentifier {
		t.Errorf("foo before a quote is %v, want identifier", got)
	}
}

func TestATripleQuotedStringCrossesLines(t *testing.T) {
	const src = "text = \"\"\"one\ntwo\nthree\"\"\"\nx = 1\n"
	spans := Highlight(src)

	// Column 8 on the opening line is inside the literal; the continued lines
	// are short, so they are asked about at their first column.
	for _, at := range []struct{ line, col int }{{0, 8}, {1, 0}, {2, 0}} {
		if class, ok := classAt(spans, at.line, at.col); !ok || class != syntax.ClassString {
			t.Errorf("line %d column %d is %v (covered: %t), want string", at.line, at.col, class, ok)
		}
	}
	if class, _ := classAt(spans, 3, 0); class != syntax.ClassIdentifier {
		t.Errorf("the line after the string is %v, want identifier — the string did not close", class)
	}
}

// A lone quote inside a triple-quoted string closes nothing: it takes three.
// Without this the closer is the same rune the opener started with, and every
// docstring that quotes anything ends in the middle of itself.
func TestALoneQuoteInsideATripleQuotedStringClosesNothing(t *testing.T) {
	const src = "text = \"\"\"say \"hi\" now\"\"\"\nx = 1\n"

	for _, inside := range []string{"hi", "now"} {
		if got := classOfFirst(t, src, inside); got != syntax.ClassString {
			t.Errorf("%q inside the docstring is %v, want string — one quote ended it", inside, got)
		}
	}
	if class, _ := classAt(Highlight(src), 1, 0); class != syntax.ClassIdentifier {
		t.Errorf("the line after it is %v, want identifier", class)
	}
}

// The same thing across a line break, which is where a docstring actually
// lives: the carried closer has to be three quotes, not one.
func TestACarriedTripleQuotedStringNeedsThreeQuotesToClose(t *testing.T) {
	const src = "text = \"\"\"first \"quoted\"\nsecond \"also\"\nthird\"\"\"\nx = 1\n"
	spans := Highlight(src)

	for line := range 3 {
		if class, ok := classAt(spans, line, 7); !ok || class != syntax.ClassString {
			t.Errorf("line %d is %v (covered: %t), want string — a lone quote closed it", line, class, ok)
		}
	}
	if class, _ := classAt(spans, 3, 0); class != syntax.ClassIdentifier {
		t.Errorf("the line after it is %v, want identifier — the string never closed", class)
	}
}

func TestOneKindOfTripleQuoteDoesNotCloseTheOther(t *testing.T) {
	const src = "text = '''one \"\"\" two'''\nx = 1\n"

	if got := classOfFirst(t, src, `"""`); got != syntax.ClassString {
		t.Errorf(`the """ inside a ''' string is %v, want string`, got)
	}
	if class, _ := classAt(Highlight(src), 1, 0); class != syntax.ClassIdentifier {
		t.Errorf("the next line is %v, want identifier — the ''' never closed", class)
	}
}

// A single-quoted string is not allowed to cross a line break, so one that
// reaches the end of a line without a backslash is coloured to there and
// dropped. Carrying it would paint the rest of the file as a string.
func TestAnUnterminatedSingleQuotedStringDoesNotCrossTheLineBreak(t *testing.T) {
	const src = "x = \"unterminated\ny = 1\n"

	if class, _ := classAt(Highlight(src), 1, 0); class != syntax.ClassIdentifier {
		t.Errorf("the line after an unterminated string is %v, want identifier", class)
	}
}

// …but a backslash at the end of the line escapes the newline, and then it
// really does carry on. That is the one case where carrying is right.
func TestABackslashAtTheEndOfALineContinuesASingleQuotedString(t *testing.T) {
	const src = "x = \"one\\\ntwo\"\ny = 1\n"

	if class, ok := classAt(Highlight(src), 1, 0); !ok || class != syntax.ClassString {
		t.Errorf("the continued line is %v (covered: %t), want string", class, ok)
	}
	if class, _ := classAt(Highlight(src), 2, 0); class != syntax.ClassIdentifier {
		t.Errorf("the line after the close is %v, want identifier", class)
	}
}

// In a raw string the backslash is kept in the value, but it still stops the
// quote after it from ending the literal — which is why rawness is not carried.
func TestABackslashEscapesTheQuoteInARawStringToo(t *testing.T) {
	const src = `p = r"\"" + "after"`

	if got := classOfFirst(t, src, `"after"`); got != syntax.ClassString {
		t.Errorf(`"after" is %v, want string — r"\"" ended one quote too early`, got)
	}
}

func TestAHashInsideAStringIsNotAComment(t *testing.T) {
	if got := classOfFirst(t, `url = "http://x/#anchor"`, "#anchor"); got != syntax.ClassString {
		t.Errorf("the # inside a string is %v, want string", got)
	}
}

// --- decorators -------------------------------------------------------------

func TestADecoratorIsAnAttribute(t *testing.T) {
	for _, src := range []string{"@property\n", "    @property\n", "@app.route\n"} {
		t.Run(src, func(t *testing.T) {
			if got := classOfFirst(t, src, "@"); got != syntax.ClassAttribute {
				t.Errorf("the decorator in %q is %v, want attribute", src, got)
			}
		})
	}
}

func TestADecoratorStopsAtItsArguments(t *testing.T) {
	const src = `@pytest.mark.parametrize("n", [1, 2])`

	if span := spanOfFirst(t, src, "@"); span.End != len("@pytest.mark.parametrize") {
		t.Errorf("the decorator is coloured to column %d, want %d", span.End, len("@pytest.mark.parametrize"))
	}
	if got := classOfFirst(t, src, `"n"`); got != syntax.ClassString {
		t.Errorf(`the "n" argument is %v, want string`, got)
	}
}

// The same rune is the matrix-multiplication operator, and only its position
// tells the two apart.
func TestAnAtSignInTheMiddleOfALineIsAnOperator(t *testing.T) {
	// With a space after it, the rune that follows already settles it. Without
	// one — `a @b` is ordinary Python — position is the only thing that does,
	// which is what this second case is for.
	for _, src := range []string{"product = a @ b", "product = a @b"} {
		t.Run(src, func(t *testing.T) {
			if got := classOfFirst(t, src, "@"); got != syntax.ClassOperator {
				t.Errorf("the @ in %q is %v, want operator", src, got)
			}
		})
	}
}

// --- the soft keywords ------------------------------------------------------

func TestMatchAndCaseAreKeywordsWhenTheyOpenABlock(t *testing.T) {
	const src = "match command.split():\n    case [\"go\", direction]:\n        pass\n"

	if got := classOfFirst(t, src, "match"); got != syntax.ClassKeyword {
		t.Errorf("match opening a statement is %v, want keyword", got)
	}
	if got := classOfFirst(t, src, "case"); got != syntax.ClassKeyword {
		t.Errorf("case opening a block is %v, want keyword", got)
	}
}

func TestMatchIsAnOrdinaryNameEverywhereElse(t *testing.T) {
	tests := []struct {
		name string
		src  string
		want syntax.Class
	}{
		{"assigned", "match = re.match(pattern, text)", syntax.ClassIdentifier},
		{"called", "if match(pattern):\n    pass\n", syntax.ClassFunction},
		{"an argument", "use(match)", syntax.ClassIdentifier},
		{"annotated", "match: str = compute()", syntax.ClassIdentifier},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			if got := classOfFirst(t, tc.src, "match"); got != tc.want {
				t.Errorf("match in %q is %v, want %v", tc.src, got, tc.want)
			}
		})
	}
}

// The boundary of the soft-keyword rule, tested rather than left to be
// discovered: a trailing comment hides the colon, and match reads as a name.
// It is the safe direction to be wrong in, and reference/languages.md says so.
func TestATrailingCommentHidesTheColonFromTheSoftKeywordRule(t *testing.T) {
	if got := classOfFirst(t, "match value:  # dispatch\n", "match"); got != syntax.ClassIdentifier {
		t.Errorf("match before a trailing comment is %v; the documented limitation says identifier", got)
	}
}

// --- what the scanner deliberately does not do ------------------------------

// An f-string's {expression} is one flat run of string, on purpose: since
// Python 3.12 it may contain anything at all, and colouring it half-properly
// breaks a format spec like "{n:{width}}".
func TestAnFStringIsNotScannedAsCodeInside(t *testing.T) {
	const src = `print(f"{count:{width}} items")`

	for _, inside := range []string{"count", "width", "items"} {
		if got := classOfFirst(t, src, inside); got != syntax.ClassString {
			t.Errorf("%q inside an f-string is %v; the whole literal is meant to be one string", inside, got)
		}
	}
}

// A docstring is a string, which is what the language calls it and what help()
// reads back. Colouring it as a comment would be a different claim, and wrong
// the moment one is assigned to a name.
func TestADocstringIsAStringAndNotAComment(t *testing.T) {
	const src = "def f():\n    \"\"\"What it does.\"\"\"\n"

	if got := classOfFirst(t, src, `"""What it does."""`); got != syntax.ClassString {
		t.Errorf("a docstring is %v, want string", got)
	}
}

// type is a builtin type as well as a soft keyword, and reads correctly as the
// type in both jobs — so it is deliberately not in isSoftKeyword.
func TestTypeIsTheBuiltinTypeInBothOfItsJobs(t *testing.T) {
	for _, src := range []string{"type(value)", "type Alias = int"} {
		if got := classOfFirst(t, src, "type"); got != syntax.ClassType {
			t.Errorf("type in %q is %v, want type", src, got)
		}
	}
}

// The walrus is an operator; every other colon is structure.
func TestTheWalrusIsAnOperatorAndAPlainColonIsNot(t *testing.T) {
	if got := classOfFirst(t, "if (n := len(text)) > 3:\n    pass\n", ":="); got != syntax.ClassOperator {
		t.Errorf(":= is %v, want operator", got)
	}
	if got := classOfFirst(t, `d = {"a": 1}`, ":"); got != syntax.ClassPunctuation {
		t.Errorf("a dict colon is %v, want punctuation", got)
	}
	if got := classOfFirst(t, "items[1:2]", ":"); got != syntax.ClassPunctuation {
		t.Errorf("a slice colon is %v, want punctuation", got)
	}
}

// --- the whole thing over a real file ---------------------------------------

func TestASweepOverRepresentativeSourceLeavesNothingUncoloured(t *testing.T) {
	// Not every rune is coloured — whitespace is not, and neither is a rune the
	// scanner steps over — but a *word* left with no span at all means the
	// dispatcher fell through, which is a defect and not a decision.
	const src = `from __future__ import annotations

import asyncio
from typing import Any


async def gather(*tasks: Any, timeout: float = 1.0) -> list[Any]:
    async with asyncio.timeout(timeout):
        return await asyncio.gather(*tasks)


class Registry(dict[str, int]):
    __slots__ = ()

    def add(self, key: str, /, *, count: int = 1) -> None:
        self[key] = self.get(key, 0) + count

    def __repr__(self) -> str:
        return f"Registry({dict(self)!r})"


lambda_ = lambda x: x if x else -x
numbers = [n**2 for n in range(10) if n % 2 == 0]
mapping = {k: v for k, v in zip("abc", [1, 2, 3])}
`

	spans := Highlight(src)
	for line, text := range strings.Split(src, "\n") {
		for col, r := range []rune(text) {
			if !syntax.IsLetter(r) && !syntax.IsDigit(r) {
				continue
			}
			if _, ok := classAt(spans, line, col); !ok {
				t.Errorf("line %d column %d (%q) is covered by no span: %q", line, col, r, text)
			}
		}
	}
}