import type { ClassBlock, CodeBlock, Inline, KindIconName, LibraryRule, Node, PropertyItem, PropertyRow, RuleCard, RuleContext, ThemeColor, TokenRole } from '../contract'
import type { Call } from './pythonTree'
import { aliasesOf, callOf, canonicalName, genericOf, lastName, stringValue, syntaxChildren, unwrapType } from './pythonTree'
import { hint, joined, name, plain, range, word } from './segments'

/**
 * SQLAlchemy 库规则：声明式模型类（有 `__tablename__`，或主体里用 `mapped_column(...)` / `Column(...)` 定义字段）整块换成「数据表」卡片：
 * 名称 / 表名 / 说明 / 继承 / 装饰器，字段、关系、索引、唯一约束各一张表，都是有才显示。
 * `__tablename__`、`__table_args__`、字段和关系的赋值（连同紧挨在它们前面的注释）并进卡片，方法等别的成员不管写在哪都照常是子块。
 * 规则只经 contract 拿信息，直接读的是 Python 语法树（tree-sitter-python 的节点类型）；用词都在词表里，按翻译插件的显示语言挑。
 */

/** 文件导入了这几个库之一才启用；Flask-SQLAlchemy 的模型写成 `db.Column(...)`，文件里可能只导入了它。 */
const LIBRARIES = ['sqlalchemy', 'flask_sqlalchemy']
const COLUMN_CALLS = new Set(['mapped_column', 'Column'])
/** 字段定义的位置参数里不是类型的那些（外键、序列、计算列……）。 */
const SCHEMA_ITEMS = new Set(['ForeignKey', 'Sequence', 'Identity', 'Computed', 'CheckConstraint', 'FetchedValue', 'DefaultClause'])
/** 写成关键字参数的类型大小：`String(length=16)`、`Numeric(precision=10, scale=2)`。 */
const SIZE_KEYWORDS = ['length', 'precision', 'scale']
const TABLE_KIND = 'sqlalchemy.table'
const TABLE_COLOR: ThemeColor = { dark: '#e8a860', light: '#a3590d' }
const TABLE_ICON: KindIconName = 'database'

/** 一种显示语言的用词。 */
interface Words {
  kind: string
  labels: Record<'name' | 'table' | 'doc' | 'bases' | 'decorators' | 'columns' | 'relationships' | 'indexes' | 'uniques', string>
  columnHeads: string[]
  relationshipHeads: string[]
  indexHeads: string[]
  count: (n: number) => string
  /** SQL 类型（SQLAlchemy 的类名；导入时起的别名已换回原名）→ 译名。 */
  types: Record<string, string>
  withTimezone: string
  enumType: string
  unannotated: string
  constraints: Record<'primaryKey' | 'unique' | 'index' | 'nullable' | 'default' | 'serverDefault' | 'foreignKey', string>
  /** 外键 ondelete / onupdate 的动作（大写）→ 译名；表里没有的用 otherOn… 拼。 */
  onDelete: Record<string, string>
  onUpdate: Record<string, string>
  otherOnDelete: (action: string) => string
  otherOnUpdate: (action: string) => string
  /** 外键后面括起来的删除、更新规则。 */
  actions: (texts: string[]) => string
  unnamed: string
  uniqueIndex: string
  relationship: Record<'backPopulates' | 'single' | 'secondary', string>
}

const ZH: Words = {
  kind: '数据表',
  labels: { name: '名称', table: '表名', doc: '说明', bases: '继承', decorators: '装饰器', columns: '字段', relationships: '关系', indexes: '索引', uniques: '唯一约束' },
  columnHeads: ['字段', '类型', '约束', '说明'],
  relationshipHeads: ['名称', '目标', '说明'],
  indexHeads: ['名称', '列'],
  count: n => `${n} 个`,
  types: {
    Integer: '整数',
    INTEGER: '整数',
    INT: '整数',
    BigInteger: '长整数',
    BIGINT: '长整数',
    SmallInteger: '短整数',
    SMALLINT: '短整数',
    String: '字符串',
    Unicode: '字符串',
    VARCHAR: '字符串',
    NVARCHAR: '字符串',
    CHAR: '字符串',
    Text: '长文本',
    UnicodeText: '长文本',
    TEXT: '长文本',
    Boolean: '布尔',
    BOOLEAN: '布尔',
    DateTime: '日期时间',
    DATETIME: '日期时间',
    TIMESTAMP: '日期时间',
    Date: '日期',
    DATE: '日期',
    Time: '时间',
    TIME: '时间',
    Interval: '时间间隔',
    Numeric: '定点数',
    NUMERIC: '定点数',
    DECIMAL: '定点数',
    Float: '浮点数',
    FLOAT: '浮点数',
    REAL: '浮点数',
    Double: '浮点数',
    DOUBLE_PRECISION: '浮点数',
    JSON: 'JSON',
    JSONB: 'JSON',
    UUID: 'UUID',
    Uuid: 'UUID',
    LargeBinary: '二进制',
    BYTEA: '二进制',
    ARRAY: '数组',
  },
  withTimezone: '（带时区）',
  enumType: '枚举',
  unannotated: '未标注',
  constraints: { primaryKey: '主键', unique: '唯一', index: '索引', nullable: '可空', default: '默认', serverDefault: '数据库默认', foreignKey: '外键' },
  onDelete: { 'CASCADE': '级联删除', 'RESTRICT': '删除时禁止', 'SET NULL': '删除时置空', 'SET DEFAULT': '删除时置默认值', 'NO ACTION': '删除时不处理' },
  onUpdate: { 'CASCADE': '级联更新', 'RESTRICT': '更新时禁止', 'SET NULL': '更新时置空', 'SET DEFAULT': '更新时置默认值', 'NO ACTION': '更新时不处理' },
  otherOnDelete: action => `删除时 ${action}`,
  otherOnUpdate: action => `更新时 ${action}`,
  actions: texts => `（${texts.join('，')}）`,
  unnamed: '未命名',
  uniqueIndex: '（唯一）',
  relationship: { backPopulates: '对方字段', single: '单个对象', secondary: '中间表' },
}

/** 显示语言 → 用词；翻译插件的显示语言这里没有时用中文。 */
const WORDS: Record<string, Words> = { zh: ZH }

/** 认一个模型类要用的：宿主给的 context、挑好的用词、别名表。 */
interface Env {
  context: RuleContext
  words: Words
  aliases: Map<string, string>
}

/** 类主体里一条给名字赋值（或只写注解）的语句；doc 是它的属性 docstring（字段的 doc、变量的 facts.doc）。 */
interface Member {
  block: CodeBlock
  name: Node
  annotation?: Node
  value?: Node
  doc?: Node
}

/** 一个字段：只写 `Mapped[...]` 注解的没有 call。 */
interface Column {
  member: Member
  call?: Call
}

interface Model {
  table?: Member
  args?: Member
  columns: Column[]
  relationships: { member: Member, call: Call }[]
}

// ─── 读语法树 ─────────────────────────────────────────────────────────────────

function isTrue(node: Node | undefined): boolean {
  return node?.type === 'true'
}

function memberOf(block: CodeBlock): Member | undefined {
  if (block.kind !== 'field' && block.kind !== 'variable' && block.kind !== 'const')
    return undefined
  const [assignment] = syntaxChildren(block.node)
  const name = assignment?.type === 'assignment' ? assignment.childForFieldName('left') : null
  if (!assignment || name?.type !== 'identifier')
    return undefined
  return {
    block,
    name,
    annotation: assignment.childForFieldName('type') ?? undefined,
    value: assignment.childForFieldName('right') ?? undefined,
    doc: 'facts' in block ? block.facts.doc : block.doc,
  }
}

/** `Mapped[X]`（`so.Mapped[X]`）里的 X；注解不是 Mapped[...] 时没有。 */
function mappedType(annotation: Node | undefined): Node | undefined {
  const generic = genericOf(annotation)
  return generic && lastName(generic.head) === 'Mapped' ? generic.args[0] : undefined
}

function isNone(node: Node): boolean {
  return unwrapType(node).type === 'none'
}

/** 类型去掉 None：`X | None`、`Optional[X]`、`Union[X, None]` → X，并说它可空；拆不出单个类型时留着整个。 */
function withoutNone(type: Node): { type: Node, optional: boolean } {
  const node = unwrapType(type)
  if (node.type === 'binary_operator' && node.childForFieldName('operator')?.type === '|') {
    const left = node.childForFieldName('left')
    const right = node.childForFieldName('right')
    if (left && right && isNone(right))
      return { type: left, optional: true }
    if (left && right && isNone(left))
      return { type: right, optional: true }
    return { type, optional: [left, right].some(side => side && withoutNone(side).optional) }
  }
  if (node.type === 'generic_type') {
    const [head, parameters] = syntaxChildren(node)
    const args = syntaxChildren(parameters)
    if (lastName(head) === 'Optional' && args.length === 1)
      return { type: args[0], optional: true }
    const rest = args.filter(arg => !isNone(arg))
    if (lastName(head) === 'Union' && rest.length < args.length)
      return { type: rest.length === 1 ? rest[0] : type, optional: true }
  }
  return { type, optional: false }
}

/** 类主体里的表名、表参数、字段和关系；既没有表名也没有 mapped_column / Column 字段的类不是模型。 */
function modelOf(block: ClassBlock, aliases: Map<string, string>): Model | undefined {
  const model: Model = { columns: [], relationships: [] }
  for (const child of block.children) {
    const member = memberOf(child)
    if (!member)
      continue
    if (member.name.text === '__tablename__') {
      model.table = member
      continue
    }
    if (member.name.text === '__table_args__') {
      model.args = member
      continue
    }
    const call = callOf(member.value, aliases)
    if (call?.name && COLUMN_CALLS.has(call.name))
      model.columns.push({ member, call })
    else if (call?.name === 'relationship')
      model.relationships.push({ member, call })
    else if (!member.value && mappedType(member.annotation))
      model.columns.push({ member })
  }
  return model.table || model.columns.some(column => column.call) ? model : undefined
}

/** `__table_args__`：元组 / 列表里的约束、索引，最后一项可以是选项字典；也可以只写一个字典。 */
function tableArgs(value: Node | undefined): { options?: Node, items: Node[] } {
  let node = value
  while (node?.type === 'parenthesized_expression')
    node = syntaxChildren(node)[0]
  if (node?.type === 'dictionary')
    return { options: node, items: [] }
  if (node?.type !== 'tuple' && node?.type !== 'list')
    return { items: [] }
  const items = syntaxChildren(node)
  const last = items.at(-1)
  return last?.type === 'dictionary' ? { options: last, items: items.slice(0, -1) } : { items }
}

function dictValue(dictionary: Node | undefined, key: string): Node | undefined {
  for (const pair of syntaxChildren(dictionary)) {
    if (pair.type === 'pair' && stringValue(pair.childForFieldName('key')) === key)
      return pair.childForFieldName('value') ?? undefined
  }
  return undefined
}

// ─── 写成属性表 ───────────────────────────────────────────────────────────────

/** 字符串字面量写内容（不带引号）、按 role 着色；别的表达式交给翻译插件翻。 */
function textOr(node: Node, role: TokenRole, env: Env): Inline[] {
  const text = stringValue(node)
  return text === undefined ? env.context.translate(node) : [{ text, role }]
}

/** 类名、枚举类这种名字：是单个名字就按类型着色、带下标，别的写法交给翻译插件。 */
function typeName(node: Node, env: Env): Inline[] {
  return node.type === 'identifier' ? [name(node, 'type')] : env.context.translate(node)
}

/** 枚举：`Enum(X)`、`SQLEnum(X)`、包着枚举类的辅助函数（`_enum(X, "…")`）写「枚举 X」；只列了取值的写「枚举 "a", "b"」。 */
function enumType(call: Call | undefined, env: Env): Inline[] {
  const out: Inline[] = [{ text: env.words.enumType, role: 'type' }]
  const [first] = call?.positional ?? []
  const values = first && first.type !== 'string' ? [first] : (call?.positional ?? []).filter(arg => arg.type === 'string')
  return values.length > 0 ? [...out, plain(' '), ...joined(values.map(value => typeName(value, env)), plain(', '))] : out
}

/** 字段的 SQL 类型：认得的换成译名，带上大小（`字符串(16)`、`定点数(10,2)`）和时区；不认得的照常翻。 */
function sqlType(node: Node, env: Env): Inline[] {
  const call = callOf(node, env.aliases)
  const original = call ? call.name : canonicalName(node, env.aliases)
  if (original && /enum/i.test(original))
    return enumType(call, env)
  const translated = original && env.words.types[original]
  if (!translated)
    return env.context.translate(node)
  const out: Inline[] = [{ text: translated, role: 'type' }]
  if (!call)
    return out
  const sizes = [...call.positional, ...SIZE_KEYWORDS.flatMap(key => call.keywords.get(key) ?? [])]
  if (sizes.length > 0)
    out.push(plain('('), ...joined(sizes.map(size => sqlType(size, env)), plain(',')), plain(')'))
  if (isTrue(call.keywords.get('timezone')))
    out.push({ text: env.words.withTimezone, role: 'type' })
  return out
}

function referentialAction(node: Node | undefined, known: Record<string, string>, other: (action: string) => string): string | undefined {
  if (!node)
    return undefined
  const action = stringValue(node)?.trim().toUpperCase()
  return action === undefined ? other(node.text) : known[action] ?? other(action)
}

/** 外键：「外键 sys_user.id（删除时禁止）」。 */
function foreignKey(call: Call, env: Env): Inline[] {
  const { words } = env
  const [target] = call.positional
  const out: Inline[] = [word(words.constraints.foreignKey)]
  if (target)
    out.push(plain(' '), ...textOr(target, 'string', env))
  const actions = [
    referentialAction(call.keywords.get('ondelete'), words.onDelete, words.otherOnDelete),
    referentialAction(call.keywords.get('onupdate'), words.onUpdate, words.otherOnUpdate),
  ].filter((text): text is string => text !== undefined)
  if (actions.length > 0)
    out.push(hint(words.actions(actions)))
  return out
}

/**
 * 可空：写了 nullable 按它；主键不可空；Column 默认可空；mapped_column 和只写注解的按 `Mapped[...]` 里有没有 None，没注解时同 Column。
 * 非空是常态，不写。
 */
function isNullable(call: Call | undefined, optional: boolean | undefined): boolean {
  const explicit = call?.keywords.get('nullable')
  if (explicit)
    return isTrue(explicit)
  if (isTrue(call?.keywords.get('primary_key')))
    return false
  return call?.name === 'Column' || optional === undefined ? true : optional
}

/** 字段表的一行：字段 | 类型 | 约束 | 说明（`comment=`，没写时用字段的属性 docstring 原文）。讲的定义是字段名。 */
function columnItem({ member, call }: Column, env: Env): PropertyItem {
  const { words, context } = env
  const mapped = mappedType(member.annotation)
  const annotated = mapped && withoutNone(mapped)
  // 第一个位置参数是字符串时是库里的列名，不是类型
  const args = (call?.positional ?? []).filter((arg, i) => !(i === 0 && arg.type === 'string'))
  const calls = args.map(arg => callOf(arg, env.aliases))
  const typeArg = args.find((_, i) => !SCHEMA_ITEMS.has(calls[i]?.name ?? ''))
  const keywords = call?.keywords ?? new Map<string, Node>()
  const constraints: Inline[][] = []
  const flag = (key: string, text: string): void => {
    if (isTrue(keywords.get(key)))
      constraints.push([word(text)])
  }
  flag('primary_key', words.constraints.primaryKey)
  flag('unique', words.constraints.unique)
  flag('index', words.constraints.index)
  if (isNullable(call, annotated ? annotated.optional : undefined))
    constraints.push([word(words.constraints.nullable)])
  for (const [key, text] of [['default', words.constraints.default], ['server_default', words.constraints.serverDefault]]) {
    const value = keywords.get(key)
    if (value)
      constraints.push([word(text), plain(' '), ...context.translate(value)])
  }
  for (const foreign of calls.filter((found): found is Call => found?.name === 'ForeignKey'))
    constraints.push(foreignKey(foreign, env))
  const comment = keywords.get('comment')
  const doc = member.doc && context.doc(member.doc)?.body
  return {
    cells: [
      [name(member.name, 'variable')],
      typeArg ? sqlType(typeArg, env) : annotated ? context.translate(annotated.type) : [hint(words.unannotated)],
      joined(constraints, hint(' · ')),
      comment ? textOr(comment, 'comment', env) : doc ? [{ text: doc.markdown, role: 'comment' }] : [],
    ],
    defines: range(member.name),
  }
}

/** 关系表的一行：名称 | 目标 | 说明（对方字段、单个对象、中间表）。目标取第一个位置参数，没写时取 `Mapped[...]` 里的类型。 */
function relationshipItem({ member, call }: { member: Member, call: Call }, env: Env): PropertyItem {
  const { words, context } = env
  const [target] = call.positional
  const mapped = mappedType(member.annotation)
  const notes: Inline[][] = []
  const back = call.keywords.get('back_populates') ?? call.keywords.get('backref')
  if (back)
    notes.push([word(words.relationship.backPopulates), plain(' '), ...textOr(back, 'variable', env)])
  if (call.keywords.get('uselist')?.type === 'false')
    notes.push([word(words.relationship.single)])
  const secondary = call.keywords.get('secondary')
  if (secondary)
    notes.push([word(words.relationship.secondary), plain(' '), ...textOr(secondary, 'string', env)])
  return {
    cells: [
      [name(member.name, 'variable')],
      target ? textOr(target, 'type', env) : mapped ? context.translate(withoutNone(mapped).type) : [],
      joined(notes, hint(' · ')),
    ],
    defines: range(member.name),
  }
}

/** 索引 / 唯一约束表的一行：名称 | 列。名称是字符串就写内容，没给写「未命名」；唯一索引的名称后面写「（唯一）」。 */
function constraintItem(nameNode: Node | undefined, columns: Node[], unique: boolean, env: Env): PropertyItem {
  const title = nameNode && nameNode.type !== 'none' ? textOr(nameNode, 'string', env) : [hint(env.words.unnamed)]
  return {
    cells: [
      unique ? [...title, hint(env.words.uniqueIndex)] : title,
      joined(columns.map(column => textOr(column, 'variable', env)), plain(', ')),
    ],
  }
}

function tableRow(label: string, heads: string[], items: PropertyItem[], env: Env): PropertyRow[] {
  return items.length > 0 ? [{ label, value: [hint(env.words.count(items.length))], columns: heads, items }] : []
}

function tableBlock(block: ClassBlock, model: Model, env: Env): RuleCard {
  const { words, context } = env
  const { facts } = block
  const { labels } = words
  const { options, items } = tableArgs(model.args?.value)
  const comment = dictValue(options, 'comment')
  const doc = (comment && context.doc(comment)?.body) ?? (facts.doc && context.doc(facts.doc)?.body)
  const indexes: PropertyItem[] = []
  const uniques: PropertyItem[] = []
  for (const item of items) {
    const call = callOf(item, env.aliases)
    if (call?.name === 'Index') {
      const [indexName, ...columns] = call.positional
      indexes.push(constraintItem(indexName, columns, isTrue(call.keywords.get('unique')), env))
    }
    else if (call?.name === 'UniqueConstraint') {
      uniques.push(constraintItem(call.keywords.get('name'), call.positional, false, env))
    }
  }
  const rows: PropertyRow[] = [
    ...(facts.name ? [{ label: labels.name, value: [name(facts.name, 'type')] }] : []),
    ...(model.table?.value ? [{ label: labels.table, value: textOr(model.table.value, 'string', env) }] : []),
    ...(doc ? [{ label: labels.doc, value: doc }] : []),
    ...(facts.bases.length > 0 ? [{ label: labels.bases, value: joined(facts.bases.map(base => context.translate(base.value)), plain(', ')) }] : []),
    ...(facts.decorators.length > 0 ? [{ label: labels.decorators, value: joined(facts.decorators.map(decorator => context.translate(decorator)), plain('\n')) }] : []),
    ...tableRow(labels.columns, words.columnHeads, model.columns.map(column => columnItem(column, env)), env),
    ...tableRow(labels.relationships, words.relationshipHeads, model.relationships.map(relationship => relationshipItem(relationship, env)), env),
    ...tableRow(labels.indexes, words.indexHeads, indexes, env),
    ...tableRow(labels.uniques, words.indexHeads, uniques, env),
  ]
  const members = [model.table, model.args, ...model.columns.map(column => column.member), ...model.relationships.map(relationship => relationship.member)]
  return {
    kind: TABLE_KIND,
    badge: { label: words.kind, color: TABLE_COLOR, icon: TABLE_ICON },
    rows,
    ...(facts.name && { defines: range(facts.name) }),
    absorbs: members.filter((member): member is Member => member !== undefined).map(member => member.block),
  }
}

export const sqlalchemyRule: LibraryRule = {
  id: 'sqlalchemy',
  label: 'SQLAlchemy 数据表',
  languageIds: ['python'],
  libraries: LIBRARIES,
  claim: (block, context) => {
    if (block.kind !== 'class')
      return undefined
    const aliases = aliasesOf(context.imports, LIBRARIES)
    const model = modelOf(block, aliases)
    return model && tableBlock(block, model, { context, words: WORDS[context.displayLanguage ?? ''] ?? ZH, aliases })
  },
}