#include "SqlContext.h"
#include "SqlLexer.h"
#include <QSet>
#include <algorithm>

namespace {

struct Tok {
	int start = 0;
	int end = 0;
	BasicTokenType type = BasicTokenType::None;
	QString text;
	QString upper; // only filled for Symbol tokens

	bool isIdent() const
	{
		return type == BasicTokenType::QuotedIdentifier
			|| (type == BasicTokenType::Symbol && !text.isEmpty() && !text[0].isDigit());
	}
	bool isSelf(QChar c) const { return type == BasicTokenType::Self && text.size() == 1 && text[0] == c; }
	bool isKeyword(const char *kw) const { return type == BasicTokenType::Symbol && upper == QLatin1String(kw); }
};

const QSet<QString> &reservedWords()
{
	static const QSet<QString> words = {
		"SELECT", "FROM", "WHERE", "GROUP", "ORDER", "HAVING", "LIMIT", "OFFSET", "FETCH", "FOR",
		"UNION", "INTERSECT", "EXCEPT", "JOIN", "INNER", "LEFT", "RIGHT", "FULL", "CROSS", "NATURAL",
		"OUTER", "ON", "USING", "SET", "RETURNING", "VALUES", "WINDOW", "AS", "LATERAL", "WITH",
		"INSERT", "INTO", "UPDATE", "DELETE", "ONLY", "TABLESAMPLE"
	};
	return words;
}

bool isReserved(const Tok &t)
{
	return t.type == BasicTokenType::Symbol && reservedWords().contains(t.upper);
}

bool isSetOperator(const Tok &t)
{
	return t.isKeyword("UNION") || t.isKeyword("INTERSECT") || t.isKeyword("EXCEPT");
}

QString unquoteIdent(const QString &text, bool &quoted)
{
	quoted = false;
	if (text.startsWith('"')) {
		quoted = true;
		QString s = text.mid(1);
		if (s.endsWith('"'))
			s.chop(1);
		s.replace(QStringLiteral("\"\""), QStringLiteral("\""));
		return s;
	}
	return text;
}

// Name as used for matching against the catalog
QString normalizedIdent(const Tok &t, bool *quotedOut = nullptr)
{
	bool quoted;
	QString s = unquoteIdent(t.text, quoted);
	if (quotedOut)
		*quotedOut = quoted;
	return quoted ? s : s.toLower();
}

std::vector<Tok> tokenize(const QString &text)
{
	std::vector<Tok> toks;
	SqlLexer lexer(text, LexerState::Null);
	for (;;) {
		SqlToken t = lexer.nextBasicToken();
		if (t.tokenType == BasicTokenType::End)
			break;
		if (t.length <= 0 && t.tokenType == BasicTokenType::None)
			break;
		if (t.tokenType == BasicTokenType::WhiteSpace || t.tokenType == BasicTokenType::NewLine)
			continue;
		Tok tok;
		tok.start = t.startPos;
		tok.end = t.startPos + t.length;
		tok.type = t.tokenType;
		tok.text = text.mid(t.startPos, t.length);
		if (tok.type == BasicTokenType::Symbol)
			tok.upper = tok.text.toUpper();
		toks.push_back(std::move(tok));
	}
	return toks;
}

bool isUnterminatedQuote(const QString &s)
{
	return s.startsWith('"') && (s.count('"') % 2 == 1);
}

bool isRelationClause(SqlClause c)
{
	return c == SqlClause::From || c == SqlClause::Join || c == SqlClause::Update || c == SqlClause::InsertInto;
}

// Relations of a parent query are visible in sub queries inside these clauses
bool parentVisible(SqlClause c)
{
	switch (c) {
	case SqlClause::Select:
	case SqlClause::Where:
	case SqlClause::On:
	case SqlClause::GroupBy:
	case SqlClause::OrderBy:
	case SqlClause::Having:
	case SqlClause::Set:
	case SqlClause::Returning:
	case SqlClause::Values:
	case SqlClause::InsertInto:
		return true;
	default:
		return false;
	}
}

class Analyzer {
public:
	using Toks = std::vector<Tok>;

	explicit Analyzer(const Toks &s) : S(s) {}

	struct Level {
		int begin = 0; // first token index inside the level
		int end = 0;   // one past last token index inside the level
		int open = -1; // index of "(" or -1 for statement level
	};

	// Direct tokens of level: nested groups are skipped, but the "(" and ")" markers are kept
	std::vector<int> directTokens(const Level &lv) const
	{
		std::vector<int> d;
		int depth = 0;
		for (int i = lv.begin; i < lv.end; ++i) {
			if (S[i].isSelf('(')) {
				if (depth == 0)
					d.push_back(i);
				++depth;
			}
			else if (S[i].isSelf(')')) {
				if (depth > 0)
					--depth;
				if (depth == 0)
					d.push_back(i);
			}
			else if (depth == 0)
				d.push_back(i);
		}
		return d;
	}

	struct Segment {
		std::vector<int> d; // direct tokens of the segment
		int cursor = 0;     // position in d of first token at or after the cursor
	};

	// The part of the level between set operators that contains token index 'cursorIdx'
	Segment segment(const Level &lv, int cursorIdx) const
	{
		std::vector<int> d = directTokens(lv);
		size_t dc = 0;
		while (dc < d.size() && d[dc] < cursorIdx)
			++dc;
		size_t b = 0;
		for (size_t i = 0; i < dc; ++i)
			if (isSetOperator(S[d[i]]))
				b = i + 1;
		size_t e = d.size();
		for (size_t i = dc; i < d.size(); ++i)
			if (isSetOperator(S[d[i]])) {
				e = i;
				break;
			}
		Segment seg;
		seg.d.assign(d.begin() + b, d.begin() + e);
		seg.cursor = static_cast<int>(dc - b);
		return seg;
	}

	SqlClause clauseAt(const Segment &seg) const
	{
		SqlClause clause = SqlClause::None;
		for (int i = 0; i < seg.cursor; ++i) {
			const Tok &t = S[seg.d[i]];
			if (t.type != BasicTokenType::Symbol)
				continue;
			const QString &k = t.upper;
			if (k == "SELECT") clause = SqlClause::Select;
			else if (k == "FROM") clause = SqlClause::From;
			else if (k == "JOIN") clause = SqlClause::Join;
			else if (k == "ON" || k == "USING") clause = SqlClause::On;
			else if (k == "WHERE") clause = SqlClause::Where;
			else if (k == "GROUP") clause = SqlClause::GroupBy;
			else if (k == "ORDER") clause = SqlClause::OrderBy;
			else if (k == "HAVING") clause = SqlClause::Having;
			else if (k == "SET") clause = SqlClause::Set;
			else if (k == "UPDATE") clause = SqlClause::Update;
			else if (k == "INTO") clause = SqlClause::InsertInto;
			else if (k == "RETURNING") clause = SqlClause::Returning;
			else if (k == "VALUES") clause = SqlClause::Values;
			else if (k == "LIMIT" || k == "OFFSET" || k == "FETCH" || k == "WINDOW" || k == "WITH")
				clause = SqlClause::None;
		}
		return clause;
	}

	// True when the token before the cursor allows a new relation name
	bool relationExpected(const Segment &seg) const
	{
		if (seg.cursor == 0)
			return true;
		const Tok &p = S[seg.d[seg.cursor - 1]];
		return p.type == BasicTokenType::Comma
			|| p.isSelf('(')
			|| p.isKeyword("FROM") || p.isKeyword("JOIN") || p.isKeyword("UPDATE")
			|| p.isKeyword("INTO") || p.isKeyword("ONLY") || p.isKeyword("LATERAL");
	}

	// Parse one relation item starting at position d in seg; returns the position after it
	int parseItem(const Segment &seg, int d, std::vector<RelationRef> &out) const
	{
		const int n = static_cast<int>(seg.d.size());
		auto tok = [&](int i) -> const Tok & { return S[seg.d[i]]; };
		while (d < n && (tok(d).isKeyword("ONLY") || tok(d).isKeyword("LATERAL")))
			++d;
		if (d >= n)
			return d;

		RelationRef ref;
		if (tok(d).isSelf('(')) {
			ref.derived = true;
			++d;
			if (d < n && tok(d).isSelf(')'))
				++d;
		}
		else if (tok(d).isIdent() && !isReserved(tok(d))) {
			std::vector<const Tok *> parts{ &tok(d) };
			++d;
			while (d + 1 < n && tok(d).isSelf('.') && tok(d + 1).isIdent()) {
				parts.push_back(&tok(d + 1));
				d += 2;
			}
			ref.name = normalizedIdent(*parts.back(), &ref.quoted);
			if (parts.size() >= 2)
				ref.schema = normalizedIdent(*parts[parts.size() - 2]);
		}
		else
			return d;

		if (d < n && tok(d).isKeyword("AS"))
			++d;
		if (d < n && tok(d).isIdent() && !isReserved(tok(d))) {
			ref.alias = normalizedIdent(tok(d));
			++d;
		}
		if (d < n && tok(d).isSelf('(') && ref.derived == false) {
			// alias column list or INSERT column list, leave to caller; skip only after an alias
			if (!ref.alias.isEmpty() && d + 1 < n && tok(d + 1).isSelf(')'))
				d += 2;
		}
		if (ref.derived) {
			if (ref.alias.isEmpty())
				return d;
			ref.name = ref.alias;
		}
		out.push_back(std::move(ref));
		return d;
	}

	void collectRelations(const Segment &seg, std::vector<RelationRef> &out) const
	{
		const int n = static_cast<int>(seg.d.size());
		bool hasSelect = false;
		for (int i = 0; i < n; ++i)
			if (S[seg.d[i]].isKeyword("SELECT"))
				hasSelect = true;

		bool inFrom = false;
		int d = 0;
		while (d < n) {
			const Tok &t = S[seg.d[d]];
			if (t.isKeyword("FROM") || t.isKeyword("JOIN") || t.isKeyword("UPDATE")) {
				inFrom = true;
				d = parseItem(seg, d + 1, out);
			}
			else if (t.isKeyword("INTO") && d > 0 && S[seg.d[d - 1]].isKeyword("INSERT")) {
				inFrom = true;
				if (!hasSelect)
					d = parseItem(seg, d + 1, out);
				else {
					std::vector<RelationRef> dummy;
					d = parseItem(seg, d + 1, dummy);
				}
			}
			else if (t.type == BasicTokenType::Comma && inFrom) {
				d = parseItem(seg, d + 1, out);
			}
			else {
				if (t.type == BasicTokenType::Symbol && reservedWords().contains(t.upper))
					inFrom = false;
				++d;
			}
		}
	}

	// Names of common table expressions defined by the WITH clause of the statement
	void collectCtes(std::vector<RelationRef> &out) const
	{
		Level top{ 0, static_cast<int>(S.size()), -1 };
		std::vector<int> d = directTokens(top);
		size_t i = 0;
		if (d.empty() || !S[d[0]].isKeyword("WITH"))
			return;
		i = 1;
		if (i < d.size() && S[d[i]].isKeyword("RECURSIVE"))
			++i;
		while (i < d.size()) {
			const Tok &nameTok = S[d[i]];
			if (!nameTok.isIdent() || isReserved(nameTok))
				return;
			RelationRef ref;
			ref.name = normalizedIdent(nameTok, &ref.quoted);
			ref.derived = true;
			++i;
			if (i < d.size() && S[d[i]].isSelf('(')) { // column list
				++i;
				if (i < d.size() && S[d[i]].isSelf(')'))
					++i;
			}
			out.push_back(std::move(ref));
			if (i >= d.size() || !S[d[i]].isKeyword("AS"))
				return;
			++i;
			while (i < d.size() && (S[d[i]].isKeyword("NOT") || S[d[i]].isKeyword("MATERIALIZED")))
				++i;
			if (i < d.size() && S[d[i]].isSelf('(')) {
				++i;
				if (i < d.size() && S[d[i]].isSelf(')'))
					++i;
			}
			else
				return;
			if (i < d.size() && S[d[i]].type == BasicTokenType::Comma)
				++i;
			else
				return;
		}
	}

	const Toks &S;
};

} // namespace

SqlContext analyzeSqlContext(const QString &text, int cursorPos)
{
	SqlContext ctx;
	cursorPos = std::clamp(cursorPos, 0, static_cast<int>(text.size()));
	ctx.prefixStart = cursorPos;

	// An unterminated quoted identifier before the cursor would swallow the remaining text
	QString work = text;
	{
		std::vector<Tok> head = tokenize(text.left(cursorPos));
		if (!head.empty() && head.back().end == cursorPos && isUnterminatedQuote(head.back().text))
			work.insert(cursorPos, QLatin1Char('"'));
	}

	std::vector<Tok> T = tokenize(work);

	// Inside comments and strings nothing is completed
	for (const Tok &t : T) {
		if (t.start < cursorPos && cursorPos <= t.end) {
			switch (t.type) {
			case BasicTokenType::Comment:
			case BasicTokenType::BlockComment:
			case BasicTokenType::OpenBlockComment:
			case BasicTokenType::QuotedString:
			case BasicTokenType::DollarQuote:
				return ctx;
			default:
				break;
			}
		}
		else if (t.type == BasicTokenType::OpenBlockComment && t.start < cursorPos)
			return ctx;
	}

	// Prefix token
	int prefIdx = -1;
	int cut = static_cast<int>(T.size());
	for (int i = 0; i < static_cast<int>(T.size()); ++i) {
		const Tok &t = T[i];
		if (t.start < cursorPos && cursorPos <= t.end && t.isIdent()) {
			prefIdx = i;
			cut = i;
			break;
		}
		if (t.start >= cursorPos) {
			cut = i;
			break;
		}
	}
	if (prefIdx >= 0) {
		const Tok &t = T[prefIdx];
		QString typed = text.mid(t.start, cursorPos - t.start);
		ctx.prefixStart = t.start;
		ctx.prefixLen = cursorPos - t.start;
		ctx.prefix = unquoteIdent(typed, ctx.prefixWasQuoted);
	}

	// Qualifier: [schema .] qualifier .
	int qualStart = cut;
	if (cut >= 2 && T[cut - 1].isSelf('.') && T[cut - 2].isIdent()) {
		ctx.qualifier = normalizedIdent(T[cut - 2]);
		qualStart = cut - 2;
		if (cut >= 4 && T[cut - 3].isSelf('.') && T[cut - 4].isIdent()) {
			ctx.qualifierSchema = normalizedIdent(T[cut - 4]);
			qualStart = cut - 4;
		}
	}
	const int removedEnd = prefIdx >= 0 ? cut + 1 : cut; // one past last removed token

	// Statement bounds
	int sb = 0;
	for (int i = qualStart - 1; i >= 0; --i)
		if (T[i].isSelf(';')) {
			sb = i + 1;
			break;
		}
	int se = static_cast<int>(T.size());
	for (int i = removedEnd; i < static_cast<int>(T.size()); ++i)
		if (T[i].isSelf(';')) {
			se = i;
			break;
		}

	std::vector<Tok> S;
	for (int i = sb; i < se; ++i)
		if (i < qualStart || i >= removedEnd)
			S.push_back(T[i]);
	const int cur = qualStart - sb; // index in S of first token after the cursor

	Analyzer an(S);

	// Enclosing paren groups
	std::vector<int> opens;
	for (int i = 0; i < cur; ++i) {
		if (S[i].isSelf('('))
			opens.push_back(i);
		else if (S[i].isSelf(')') && !opens.empty())
			opens.pop_back();
	}
	auto levelFor = [&](int open) {
		Analyzer::Level lv;
		lv.open = open;
		lv.begin = open + 1;
		lv.end = static_cast<int>(S.size());
		if (open >= 0) {
			int depth = 1;
			for (int i = open + 1; i < static_cast<int>(S.size()); ++i) {
				if (S[i].isSelf('(')) ++depth;
				else if (S[i].isSelf(')') && --depth == 0) {
					lv.end = i;
					break;
				}
			}
		}
		else
			lv.begin = 0;
		return lv;
	};

	// Levels from innermost to outermost
	std::vector<Analyzer::Level> levels;
	for (int i = static_cast<int>(opens.size()) - 1; i >= 0; --i)
		levels.push_back(levelFor(opens[i]));
	levels.push_back(levelFor(-1));

	SqlClause clause = SqlClause::None;
	bool relExpected = false;
	bool clauseResolved = false;
	bool visible = true;
	for (size_t li = 0; li < levels.size(); ++li) {
		const Analyzer::Level &lv = levels[li];
		const int cursorInLevel = li == 0 ? cur : levels[li - 1].open;
		Analyzer::Segment seg = an.segment(lv, cursorInLevel);
		const SqlClause lc = an.clauseAt(seg);

		// a level without clause keyword is a parenthesised expression and inherits the clause of its parent
		if (!clauseResolved && lc != SqlClause::None) {
			clause = lc;
			relExpected = isRelationClause(lc) && an.relationExpected(seg);
			clauseResolved = true;
		}
		if (li >= 1)
			visible = visible && parentVisible(lc);
		if (visible)
			an.collectRelations(seg, ctx.relations);
	}
	an.collectCtes(ctx.relations);

	ctx.clause = clause;
	const bool hasQualifier = !ctx.qualifier.isEmpty();
	switch (clause) {
	case SqlClause::Select:
	case SqlClause::On:
	case SqlClause::Where:
	case SqlClause::GroupBy:
	case SqlClause::OrderBy:
	case SqlClause::Having:
	case SqlClause::Set:
	case SqlClause::Returning:
	case SqlClause::Values:
		ctx.kind = hasQualifier ? CompletionKind::QualifiedColumn : CompletionKind::Column;
		break;
	case SqlClause::From:
	case SqlClause::Join:
	case SqlClause::Update:
	case SqlClause::InsertInto:
		if (hasQualifier)
			ctx.kind = CompletionKind::Relation;
		else if (relExpected)
			ctx.kind = clause == SqlClause::Join ? CompletionKind::JoinTarget : CompletionKind::Relation;
		else if (clause == SqlClause::InsertInto && !levels.empty() && levels[0].open >= 0)
			ctx.kind = CompletionKind::Column;
		break;
	default:
		break;
	}
	return ctx;
}