"""PubMed 查询语法解析器

将用户输入的 PubMed 风格查询（如 `"lung cancer"[TI] AND "pembrolizumab"[TIAB]`）
解析为结构化数据，供 AdvancedSearchEngine.search() 消费。

支持的语法：
  - 字段标签 (case-insensitive): [TI] [AB] [TIAB] [AU] [TA] [MH] [MAJR] [PT] [DP] [PMID] [DOI]
  - 布尔运算符: AND / OR / NOT（大小写不敏感，前后需词边界）
  - 引号短语: "exact phrase"
  - 括号分组: (a OR b) AND c
  - 日期范围: 2024:2026[DP]

设计原则：
  - 向后兼容：不含特殊语法的纯文本查询原样返回为 plain_terms
  - 容错：语法错误时退化到 plain_terms，不抛异常
  - 组合：解析结果映射到 AdvancedSearchEngine.search() 的已有参数
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from enum import Enum, auto


# ─── 字段标签映射 ───

_FIELD_TAG_MAP: dict[str, str] = {
    "TI": "title",
    "AB": "abstract",
    "TIAB": "all",       # 映射到 field=all（title+abstract covered）
    "AU": "author",
    "TA": "journal",
}

# MeSH/PT/DP/PMID/DOI 需要特殊处理，不直接映射到 field 参数
_SPECIAL_FIELDS = {"MH", "MAJR", "PT", "DP", "PMID", "DOI"}

# 所有合法字段标签
_ALL_FIELD_TAGS = set(_FIELD_TAG_MAP.keys()) | _SPECIAL_FIELDS


# ─── Token Types ───

class TokenType(Enum):
    QUOTED = auto()       # "text in quotes"
    FIELD = auto()        # [TI], [MH], etc.
    AND = auto()          # AND operator
    OR = auto()           # OR operator
    NOT = auto()          # NOT operator
    LPAREN = auto()       # (
    RPAREN = auto()       # )
    COLON = auto()        # :
    WORD = auto()         # unquoted text
    NUMBER = auto()       # digits (for PMID, year)
    EOF = auto()          # end of input


@dataclass
class Token:
    type: TokenType
    value: str = ""


# ─── Tokeniser ───

_TOKEN_PATTERNS: list[tuple[TokenType, str]] = [
    (TokenType.QUOTED, r'"(?:[^"\\]|\\.)*"'),
    (TokenType.FIELD, r'\[(?:' + '|'.join(_ALL_FIELD_TAGS) + r')\]'),
    (TokenType.AND, r'\bAND\b'),
    (TokenType.OR, r'\bOR\b'),
    (TokenType.NOT, r'\bNOT\b'),
    (TokenType.LPAREN, r'\('),
    (TokenType.RPAREN, r'\)'),
    (TokenType.COLON, r':'),
    (TokenType.NUMBER, r'\d+'),
    (TokenType.WORD, r'[^\s"\[\]():]+'),
]

_TOKEN_RE = re.compile(
    '|'.join(f'(?P<{t.name}>{p})' for t, p in _TOKEN_PATTERNS),
    re.IGNORECASE,
)


def tokenise(query: str) -> list[Token]:
    """将 PubMed 查询字符串分片为 Token 列表。"""
    tokens: list[Token] = []
    for m in _TOKEN_RE.finditer(query):
        for name, value in m.groupdict().items():
            if value is not None:
                ttype = TokenType[name]
                tokens.append(Token(ttype, value))
    tokens.append(Token(TokenType.EOF))
    return tokens


# ─── AST Nodes ───

@dataclass
class Term:
    """单个搜索词项。"""
    text: str
    exact: bool = False          # True if quoted phrase
    field: str | None = None     # None = plain (no field tag); "title"/"abstract"/etc for mapped; "MH"/"PT"/etc for special
    is_not: bool = False         # True if preceded by NOT


@dataclass
class ParsedPubmedQuery:
    """PubMed 查询解析后的结构化表示。"""
    title_terms: list[Term] = field(default_factory=list)          # [TI]
    abstract_terms: list[Term] = field(default_factory=list)       # [AB]
    tiab_terms: list[Term] = field(default_factory=list)           # [TIAB]
    author_terms: list[Term] = field(default_factory=list)         # [AU]
    journal_terms: list[Term] = field(default_factory=list)        # [TA]
    mesh_terms: list[str] = field(default_factory=list)            # [MH]
    majr_terms: list[str] = field(default_factory=list)            # [MAJR]
    pub_types: list[str] = field(default_factory=list)             # [PT]
    doi_terms: list[str] = field(default_factory=list)             # [DOI]
    pmid_terms: list[int] = field(default_factory=list)            # [PMID]
    date_from: str | None = None                                   # [DP] lower bound (YYYY-MM-DD)
    date_to: str | None = None                                     # [DP] upper bound (YYYY-MM-DD)
    year_from: int | None = None                                   # [DP] year lower
    year_to: int | None = None                                     # [DP] year upper
    plain_terms: list[Term] = field(default_factory=list)          # no field tag
    boolean_operator: str = "and"                                  # "and" | "or" | "mixed"
    has_not: bool = False                                          # contains NOT
    not_terms: list[Term] = field(default_factory=list)            # terms under NOT


# ─── Parser ───

class ParseError(ValueError):
    """查询语法解析错误（调用方应退化到 plain text 搜索）。"""


class PubmedQueryParser:
    """Recursive descent parser for PubMed query syntax.

    Grammar:
      query     → or_expr
      or_expr   → and_expr (OR and_expr)*
      and_expr  → not_expr (AND not_expr)*
      not_expr  → NOT not_expr | primary
      primary   → atom FIELD? | LPAREN query RPAREN
      atom      → QUOTED | WORD | number (COLON number)?
    """

    def __init__(self, tokens: list[Token]):
        self.tokens = tokens
        self.pos = 0

    def peek(self) -> Token:
        return self.tokens[self.pos]

    def peek_n(self, n: int) -> Token | None:
        idx = self.pos + n
        return self.tokens[idx] if idx < len(self.tokens) else None

    def advance(self) -> Token:
        t = self.tokens[self.pos]
        self.pos += 1
        return t

    def expect(self, *types: TokenType) -> Token:
        t = self.peek()
        if t.type not in types:
            raise ParseError(
                f"Expected one of {[tt.name for tt in types]}, got {t.type.name}({t.value!r})"
            )
        return self.advance()

    def parse(self) -> ParsedPubmedQuery:
        """入口：解析完整的查询字符串。"""
        result = ParsedPubmedQuery()
        try:
            terms = self._parse_or_expr(result)
        except ParseError:
            return ParsedPubmedQuery()

        # Detect boolean operator from token stream
        has_and = any(t.type == TokenType.AND for t in self.tokens)
        has_or = any(t.type == TokenType.OR for t in self.tokens)
        if has_and and has_or:
            result.boolean_operator = "mixed"
        elif has_or and not has_and:
            result.boolean_operator = "or"

        result.has_not = any(t.is_not for t in terms if not getattr(t, '_is_range_end', False))
        result.not_terms = [t for t in terms if t.is_not and not getattr(t, '_is_range_end', False)]

        for t in terms:
            if getattr(t, '_is_range_end', False):
                continue
            self._dispatch_term(result, t)

        return result

    def _dispatch_term(self, result: ParsedPubmedQuery, term: Term) -> None:
        if term.field == "title":
            result.title_terms.append(term)
        elif term.field == "abstract":
            result.abstract_terms.append(term)
        elif term.field == "all":  # TIAB
            result.tiab_terms.append(term)
        elif term.field == "author":
            result.author_terms.append(term)
        elif term.field == "journal":
            result.journal_terms.append(term)
        elif term.field == "MH":
            result.mesh_terms.append(term.text)
        elif term.field == "MAJR":
            result.majr_terms.append(term.text)
        elif term.field == "PT":
            result.pub_types.append(term.text)
        elif term.field == "PMID":
            try:
                result.pmid_terms.append(int(term.text))
            except ValueError:
                result.plain_terms.append(term)
        elif term.field == "DOI":
            result.doi_terms.append(term.text)
        elif term.field is None:
            result.plain_terms.append(term)
        elif term.field == "__RANGE_DP__":
            # Already handled inline during parse; skip.
            pass
        else:
            result.plain_terms.append(term)

    def _parse_or_expr(self, result: ParsedPubmedQuery) -> list[Term]:
        """or_expr → and_expr (OR and_expr)*"""
        left = self._parse_and_expr(result)
        while self.peek().type == TokenType.OR:
            self.advance()
            right = self._parse_and_expr(result)
            left.extend(right)
        return left

    def _is_primary_start(self, token: Token) -> bool:
        """Check if token could start a primary expression."""
        return token.type in (TokenType.WORD, TokenType.QUOTED, TokenType.NUMBER,
                              TokenType.LPAREN, TokenType.NOT)

    def _parse_and_expr(self, result: ParsedPubmedQuery) -> list[Term]:
        """and_expr → not_expr (AND? not_expr)*

        AND is optional between consecutive primaries (implicit AND).
        e.g. ``smith j[AU]`` → smith AND j[AU]
        """
        left = self._parse_not_expr(result)
        while True:
            tok = self.peek()
            if tok.type == TokenType.AND:
                self.advance()
            elif self._is_primary_start(tok):
                pass  # implicit AND — continue without consuming
            else:
                break
            right = self._parse_not_expr(result)
            left.extend(right)
        return left

    def _parse_not_expr(self, result: ParsedPubmedQuery) -> list[Term]:
        """not_expr → NOT not_expr | primary"""
        if self.peek().type == TokenType.NOT:
            self.advance()
            return self._parse_primary(result, negated=True)
        return self._parse_primary(result, negated=False)

    def _parse_primary(self, result: ParsedPubmedQuery, negated: bool = False) -> list[Term]:
        """primary → atom FIELD? | LPAREN query RPAREN"""
        if self.peek().type == TokenType.LPAREN:
            self.advance()
            terms = self._parse_or_expr(result)
            self.expect(TokenType.RPAREN)
            if negated:
                for t in terms:
                    t.is_not = True
            return terms

        return self._parse_atom(result, negated)

    def _parse_atom(self, result: ParsedPubmedQuery, negated: bool = False) -> list[Term]:
        """Parse a single atom with optional field tag.

        Handles:
          - "quoted phrase"[FIELD]
          - word[FIELD]
          - NUMBER:NUMBER[DP]  (year/date range)
          - NUMBER  (bare number)
          - word    (bare word)
        """
        # Look ahead for NUMBER:NUMBER pattern
        t0 = self.peek()
        t1 = self.peek_n(1)
        t2 = self.peek_n(2)

        if (t0.type == TokenType.NUMBER
                and t1 is not None and t1.type == TokenType.COLON
                and t2 is not None and t2.type == TokenType.NUMBER):
            return self._parse_range(result, negated)

        # Normal atom
        token = self.advance()
        text = token.value.strip('"') if token.type == TokenType.QUOTED else token.value
        is_exact = (token.type == TokenType.QUOTED)
        field = None
        if self.peek().type == TokenType.FIELD:
            ft = self.advance()
            field = ft.value[1:-1].upper()
            if field in _FIELD_TAG_MAP:
                field = _FIELD_TAG_MAP[field]

        return [Term(text, exact=is_exact, field=field, is_not=negated)]

    def _parse_range(self, result: ParsedPubmedQuery, negated: bool = False) -> list[Term]:
        """Parse NUMBER:NUMBER[FIELD] — handles DP (year) ranges specially."""
        start_val = self.advance().value  # NUMBER
        self.advance()                     # COLON
        end_val = self.advance().value     # NUMBER

        field = None
        if self.peek().type == TokenType.FIELD:
            ft = self.advance()
            field = ft.value[1:-1].upper()

        if field == "DP":
            try:
                sy, ey = int(start_val), int(end_val)
            except ValueError:
                sy = ey = None
            # Year-only range (e.g., 2024:2026[DP])
            if start_val.isdigit() and len(start_val) == 4:
                result.year_from = sy
                result.year_to = ey
            else:
                # Full date range (e.g., 2024/01/01:2024/12/31[DP])
                result.date_from = start_val
                result.date_to = end_val
            # Return a marker term so not-terms tracking knows about it, but _dispatch skips.
            marker = Term(f"{start_val}:{end_val}", field="__RANGE_DP__", is_not=negated)
            marker._is_range_end = True
            return [marker]

        # Non-DP range or no field → plain text
        txt = f"{start_val}:{end_val}"
        if field:
            txt = f"{txt}[{field}]"
        return [Term(txt, field=field, is_not=negated)]


# ─── Public API ───

def is_pubmed_syntax(query: str) -> bool:
    """快速检测查询串是否包含 PubMed 语法特征。

    检测条件（任意满足其一）：
    1. 包含 `[...]`（字段标签）
    2. 包含布尔运算符 AND/OR/NOT（大写，词边界）
    """
    if not query or not query.strip():
        return False
    if re.search(r'\[(' + '|'.join(_ALL_FIELD_TAGS) + r')\]', query, re.IGNORECASE):
        return True
    if re.search(r'\b(AND|OR|NOT)\b', query):
        return True
    return False


def parse_pubmed_query(query: str) -> ParsedPubmedQuery:
    """解析 PubMed 查询字符串。

    如果解析失败（非法语法），返回全部内容为 plain_terms 的空结构。
    """
    if not query or not query.strip():
        return ParsedPubmedQuery()

    try:
        tokens = tokenise(query)
        parser = PubmedQueryParser(tokens)
        return parser.parse()
    except (ParseError, IndexError, ValueError):
        return ParsedPubmedQuery()


def extract_pubmed_query_for_prisma(query: str) -> tuple[str, list[str]]:
    """提取 PubMed 查询和所用 MeSH 词，供 PRISMA 导出。

    Returns:
        (normalized_query, mesh_terms_used)
    """
    parsed = parse_pubmed_query(query)
    mesh_used = list(set(parsed.mesh_terms + parsed.majr_terms))
    mesh_used.sort()

    # 标准化：统一字段大写
    normalized = re.sub(
        r'\[(\w+)\]',
        lambda m: f'[{m.group(1).upper()}]',
        query,
    )
    return normalized, mesh_used
