import type { Access, ClassBlock, CodeBlock, ImportFacts, IndexEntry, Inline, KindIconName, LibraryRule, Node, PropertyItem, PropertyRow, RuleCard, RuleContext, ThemeColor, WorkspaceIndex } from '../contract'
import type { Call, LibraryNames } from './pythonTree'
import { annotatedAssignments, memberDocs, withLeadingComments } from './children'
import { aliasesOf, callOf, canonicalName, genericOf, libraryName, libraryNames, stringValue, syntaxChildren, unwrapType } from './pythonTree'
import { hint, joined, name, plain, range, word } from './segments'

/**
 * Pydantic 的 `BaseModel` 子类整块换成「数据模型」卡片：名称 / 说明 / 访问级别 / 继承 / 选项 / 字段表，有才显示。
 * 父类链上有 pydantic 的 BaseModel（含别名、`pydantic.BaseModel`），或同一文件里继承另一个模型，都算。
 * 父类是从别的模块导入的名字时，问工作区索引那个类是不是模型：还没查到就先当普通类，查到再换成这张卡片。
 * 字段是类主体里带注解的赋值，表和数据类同一套列；`model_config = ConfigDict(...)` 是选项，不是字段。
 * 方法、`@field_validator`、`ClassVar`、没注解的类属性照常是子块。
 */

const LIBRARY = 'pydantic'
/** ClassVar 从这几个库导入时的别名要换回原名。 */
const TYPE_LIBRARIES = ['typing', 'typing_extensions']
const KIND = 'pydantic.model'
const COLOR: ThemeColor = { dark: '#c678dd', light: '#8b3db5' }
const ICON: KindIconName = 'braces'
/** `Field(...)` 里讲默认值、说明的参数，不进标记。 */
const FIELD_VALUES = new Set(['default', 'default_factory', 'description'])

/** 一种约束的说法：值前面的词、值后面的词，有才写。 */
interface Phrase { before?: string, after?: string }

/** 一种显示语言的用词。 */
interface Words {
  kind: string
  labels: Record<'name' | 'doc' | 'access' | 'bases' | 'options' | 'fields', string>
  heads: string[]
  count: (n: number) => string
  access: Record<Access, string>
  /** `default_factory=f`：「每次新建 f()」。 */
  factory: string
  /** `ConfigDict(名=True/False)` 的译法，按「名=True」「名=False」查；没有的照写。 */
  options: Record<string, string>
  /** `extra="forbid"` 这类字符串值。 */
  extra: Record<string, string>
  /** `Field(...)` 里常见约束的说法；没有的照写参数名和值。 */
  constraints: Record<string, Phrase>
}

const ZH: Words = {
  kind: '数据模型',
  labels: { name: '名称', doc: '说明', access: '访问级别', bases: '继承', options: '选项', fields: '字段' },
  heads: ['字段', '类型', '默认值', '说明'],
  count: n => `${n} 个`,
  access: { public: '公共', private: '私有', protected: '受保护' },
  factory: '每次新建',
  options: {
    'strict=True': '严格',
    'strict=False': '不严格',
    'frozen=True': '不可修改',
    'frozen=False': '可修改',
    'from_attributes=True': '可从对象读取',
    'from_attributes=False': '不从对象读取',
    'populate_by_name=True': '可用字段名填',
    'validate_assignment=True': '赋值时也校验',
    'arbitrary_types_allowed=True': '允许任意类型',
    'use_enum_values=True': '用枚举的值',
    'str_strip_whitespace=True': '去掉首尾空白',
    'validate_default=True': '也校验默认值',
  },
  extra: { forbid: '禁止多余字段', allow: '允许多余字段', ignore: '忽略多余字段' },
  constraints: {
    min_length: { before: '至少', after: ' 个字符' },
    max_length: { before: '最长' },
    ge: { before: '≥' },
    le: { before: '≤' },
    gt: { before: '>' },
    lt: { before: '<' },
    multiple_of: { after: ' 的倍数' },
    pattern: { before: '匹配' },
    min_items: { before: '至少', after: ' 项' },
    max_items: { before: '最多', after: ' 项' },
    max_digits: { before: '最多', after: ' 位数字' },
    decimal_places: { before: '小数', after: ' 位' },
  },
}

const WORDS: Record<string, Words> = { 'zh-cn': ZH }

interface Env {
  context: RuleContext
  words: Words
  /** pydantic 里的 BaseModel、Field、ConfigDict，含导入时的别名。 */
  names: LibraryNames
  /** 注解里 ClassVar 导入时起的别名 → 原名。 */
  aliases: Map<string, string>
}

interface Field {
  block: CodeBlock
  name: Node
  annotation: Node
  value?: Node
}

/** 模块名是 pydantic，或是它的子模块（`pydantic.main`）。 */
function fromPydantic(module: string): boolean {
  return module === LIBRARY || module.startsWith(`${LIBRARY}.`) || module.startsWith(`${LIBRARY}/`)
}

/** pydantic 以及它的子模块（`pydantic.main`）里导入的名字。 */
function pydanticNames(imports: readonly ImportFacts[]): LibraryNames {
  const names = libraryNames(imports, LIBRARY)
  for (const facts of imports) {
    if (!facts.sources.some(source => fromPydantic(source.text) && source.text !== LIBRARY))
      continue
    for (const item of facts.items) {
      if (item.kind === 'named')
        names.imported.set(item.alias?.text ?? item.name.text, item.name.text)
    }
  }
  return names
}

/** 模块顶层的类：名字和父类。装饰器包着的类也算。 */
function moduleClasses(root: Node): { name: Node, bases: Node[] }[] {
  const classes: { name: Node, bases: Node[] }[] = []
  for (const statement of syntaxChildren(root)) {
    const definition = statement.type === 'decorated_definition' ? statement.childForFieldName('definition') : statement
    if (definition?.type !== 'class_definition')
      continue
    const className = definition.childForFieldName('name')
    if (className)
      classes.push({ name: className, bases: syntaxChildren(definition.childForFieldName('superclasses')) })
  }
  return classes
}

/**
 * 这一文件里能确定是模型的类名：父类是 pydantic 的 BaseModel，或是本文件里另一个模型。
 * 父类是从别的模块导入的名字时这里不算，那要等工作区索引。
 */
function localModels(classes: readonly { name: Node, bases: Node[] }[], names: LibraryNames): Set<string> {
  const models = new Set<string>()
  let grew = true
  while (grew) {
    grew = false
    for (const cls of classes) {
      if (models.has(cls.name.text))
        continue
      const hit = cls.bases.some(base => libraryName(base, names) === 'BaseModel' || (base.type === 'identifier' && models.has(base.text)))
      if (hit) {
        models.add(cls.name.text)
        grew = true
      }
    }
  }
  return models
}

/**
 * 从语法树认 pydantic 的导入。索引建的时候没有翻译插件解析好的导入表，只能自己看导入语句。
 * `import pydantic`、`from pydantic import BaseModel as BM`、`from pydantic.xxx import ...` 都算。
 */
function namesInSource(root: Node): LibraryNames {
  const names: LibraryNames = { modules: new Set(), imported: new Map() }
  for (const statement of syntaxChildren(root)) {
    if (statement.type === 'import_statement') {
      for (const item of statement.childrenForFieldName('name')) {
        const original = item.type === 'aliased_import' ? item.childForFieldName('name') : item
        const alias = item.type === 'aliased_import' ? item.childForFieldName('alias') : null
        if (original && fromPydantic(original.text))
          names.modules.add(alias?.text ?? original.text)
      }
    }
    if (statement.type !== 'import_from_statement')
      continue
    const module = statement.childForFieldName('module_name')
    if (!module || !fromPydantic(module.text))
      continue
    for (const item of statement.childrenForFieldName('name')) {
      const original = item.type === 'aliased_import' ? item.childForFieldName('name') : item
      const alias = item.type === 'aliased_import' ? item.childForFieldName('alias') : null
      if (original)
        names.imported.set(alias?.text ?? original.text, original.text)
    }
  }
  return names
}

/** 这一文件里本地就能确定的模型，记成索引项：别的文件继承这个名字时来查。 */
function modelEntries(root: Node): IndexEntry[] {
  const classes = moduleClasses(root)
  const models = localModels(classes, namesInSource(root))
  return classes.flatMap(cls => models.has(cls.name.text) ? [{ key: cls.name.text, name: range(cls.name), value: 'model' }] : [])
}

/** 哪些类是模型。只扫源码里写了 BaseModel 的文件，子类文件不在这里解析。 */
export const MODEL_INDEX: WorkspaceIndex = {
  id: 'pydantic.models',
  languageId: 'python',
  include: '**/*.py',
  mentions: source => /\bBaseModel\b/.test(source),
  entries: root => modelEntries(root),
}

/**
 * 是模型就认领。父类是别的文件里的名字时问索引：查到是模型才认，没查到先不认（宿主稍后重排）。
 */
function isModel(block: ClassBlock, env: Env): boolean {
  const className = block.facts.name?.text
  if (!className)
    return false
  const classes = moduleClasses(block.node.tree.rootNode)
  if (localModels(classes, env.names).has(className))
    return true
  const cls = classes.find(item => item.name.text === className)
  if (!cls)
    return false
  return cls.bases.some((base) => {
    if (base.type !== 'identifier' || classes.some(item => item.name.text === base.text))
      return false
    return env.context.lookup(MODEL_INDEX, base)?.value === 'model'
  })
}

/** 注解最外层的名字（`ClassVar[int]` → ClassVar），导入时的别名换回原名。 */
function annotationName(annotation: Node, env: Env): string | undefined {
  return canonicalName(genericOf(annotation)?.head ?? unwrapType(annotation), env.aliases)
}

/** 调用的是 pydantic 里的这个名字（`Field(...)`、`ConfigDict(...)`，含别名）时给出实参。 */
function namedCall(value: Node | null | undefined, expected: string, names: LibraryNames): Call | undefined {
  const callee = value?.type === 'call' ? value.childForFieldName('function') : null
  return libraryName(callee, names) === expected ? callOf(value, new Map()) : undefined
}

function isEllipsis(node: Node): boolean {
  return node.type === 'ellipsis'
}

/** 类主体里的字段：带注解的赋值，去掉 ClassVar 和 `model_config`。 */
function fieldsOf(children: readonly CodeBlock[], env: Env): Field[] {
  return annotatedAssignments(children).filter(field => field.name.text !== 'model_config' && annotationName(field.annotation, env) !== 'ClassVar')
}

/** `model_config = ...`：没注解的类属性，或写成带注解的赋值。不是字段。 */
function configOf(children: readonly CodeBlock[]): { block: CodeBlock, value?: Node } | undefined {
  for (const block of children) {
    if (block.kind === 'variable' || block.kind === 'const') {
      const [declarator, ...rest] = block.facts.declarators
      if (rest.length === 0 && declarator?.name.type === 'identifier' && declarator.name.text === 'model_config')
        return { block, value: declarator.value }
    }
    const [field] = block.kind === 'field' ? annotatedAssignments([block]) : []
    if (field?.name.text === 'model_config')
      return { block, value: field.value }
  }
  return undefined
}

/** `ConfigDict(...)` 里写了的项；不是这个调用就整段照译。没写的不出现。 */
function optionsText(value: Node | undefined, env: Env): Inline[] {
  const call = namedCall(value, 'ConfigDict', env.names)
  if (!call)
    return value ? env.context.translate(value) : []
  const parts = [...call.keywords].map(([key, arg]) => optionMark(key, arg, env))
  return joined([
    ...parts,
    ...call.positional.map(arg => env.context.translate(arg)),
    ...call.splats.map(arg => [plain('**'), ...env.context.translate(arg)]),
  ], hint(' · '))
}

function optionMark(key: string, value: Node, env: Env): Inline[] {
  const flag = value.type === 'true' ? 'True' : value.type === 'false' ? 'False' : undefined
  const known = flag ? env.words.options[`${key}=${flag}`] : undefined
  if (known)
    return [word(known)]
  if (key === 'extra') {
    const content = stringValue(value)
    const text = content ? env.words.extra[content] : undefined
    if (text)
      return [word(text)]
  }
  return [plain(`${key}=`), ...env.context.translate(value)]
}

/** `Field(...)` 的默认值：关键字 `default=` 优先，否则第一个位置参数；`...` 表示必填，不算默认值。 */
function defaultNode(call: Call): Node | undefined {
  const keyword = call.keywords.get('default')
  if (keyword)
    return isEllipsis(keyword) ? undefined : keyword
  const [positional] = call.positional
  return positional && !isEllipsis(positional) ? positional : undefined
}

/** `default_factory=f`：「每次新建 f()」；没参数的 lambda 直接写它的结果。 */
function factoryText(factory: Node, env: Env): Inline[] {
  const body = factory.type === 'lambda' && !factory.childForFieldName('parameters') ? factory.childForFieldName('body') : null
  return [word(env.words.factory), plain(' '), ...(body ? env.context.translate(body) : [...env.context.translate(factory), plain('()')])]
}

function constraintMark(key: string, value: Node, env: Env): Inline[] {
  const phrase = env.words.constraints[key]
  if (!phrase)
    return [plain(`${key}=`), ...env.context.translate(value)]
  return [
    ...(phrase.before ? [word(phrase.before), plain(' ')] : []),
    ...env.context.translate(value),
    ...(phrase.after ? [plain(phrase.after)] : []),
  ]
}

/** 默认值那一格：字面量、`Field(default=...)`、`每次新建 f()`，后面接约束标记。 */
function defaultCell(field: Field, call: Call | undefined, env: Env): Inline[] {
  const parts: Inline[][] = []
  if (!call && field.value)
    parts.push(env.context.translate(field.value))
  const value = call && defaultNode(call)
  if (value)
    parts.push(env.context.translate(value))
  const factory = call?.keywords.get('default_factory')
  if (factory)
    parts.push(factoryText(factory, env))
  for (const [key, arg] of call?.keywords ?? []) {
    if (!FIELD_VALUES.has(key))
      parts.push(constraintMark(key, arg, env))
  }
  return joined(parts, hint(' · '))
}

/** 说明：`Field(description=...)` 优先；没有就用字段上方注释、行尾注释、属性 docstring。 */
function docCell(field: Field, children: readonly CodeBlock[], call: Call | undefined, env: Env): Inline[] {
  const description = call?.keywords.get('description')
  const literal = stringValue(description)
  if (literal)
    return [{ text: literal, role: 'comment' }]
  if (description)
    return env.context.translate(description)
  const text = memberDocs(children, field.block).map(node => env.context.doc(node)?.markdown).filter(Boolean).join(' ')
  return text ? [{ text, role: 'comment' }] : []
}

function fieldItem(field: Field, children: readonly CodeBlock[], env: Env): PropertyItem {
  const call = namedCall(field.value, 'Field', env.names)
  return {
    cells: [
      [name(field.name, 'variable')],
      env.context.translate(field.annotation),
      defaultCell(field, call, env),
      docCell(field, children, call, env),
    ],
    defines: range(field.name),
  }
}

function modelCard(block: ClassBlock, env: Env): RuleCard {
  const { words, context } = env
  const { labels } = words
  const { facts, children } = block
  const config = configOf(children)
  const fields = fieldsOf(children, env)
  const doc = facts.doc && context.doc(facts.doc)
  const options = optionsText(config?.value, env)
  const rows: PropertyRow[] = [
    ...(facts.name ? [{ label: labels.name, value: [name(facts.name, 'type')] }] : []),
    ...(doc ? [{ label: labels.doc, value: doc }] : []),
    ...(facts.access ? [{ label: labels.access, value: [{ text: words.access[facts.access], role: 'keyword' as const }] }] : []),
    ...(facts.bases.length > 0 ? [{ label: labels.bases, value: joined(facts.bases.map(base => context.translate(base.value)), plain(', ')) }] : []),
    ...(options.length > 0 ? [{ label: labels.options, value: options }] : []),
    ...(fields.length > 0
      ? [{
          label: labels.fields,
          value: [hint(words.count(fields.length))],
          columns: words.heads,
          items: fields.map(field => fieldItem(field, children, env)),
        }]
      : []),
  ]
  return {
    kind: KIND,
    badge: { label: words.kind, color: COLOR, icon: ICON },
    rows,
    ...(facts.name && { defines: range(facts.name) }),
    absorbs: withLeadingComments(children, [...fields.map(field => field.block), ...(config ? [config.block] : [])]),
  }
}

export const pydanticRule: LibraryRule = {
  id: LIBRARY,
  label: 'Pydantic 数据模型',
  languageIds: ['python'],
  // 子类文件常常只导入父模型、不再导入 pydantic，不能靠这个导入来决定看不看
  libraries: [],
  claim: (block, context) => {
    if (block.kind !== 'class')
      return undefined
    const env: Env = {
      context,
      words: WORDS[context.displayLanguage ?? ''] ?? ZH,
      names: pydanticNames(context.imports),
      aliases: aliasesOf(context.imports, TYPE_LIBRARIES),
    }
    return isModel(block, env) ? modelCard(block, env) : undefined
  },
}