import React, { useState, useImperativeHandle, forwardRef, useEffect } from 'react';
import { roleApi, RoleVO } from '../api/Role';
import { menuApi, MenuVO } from '../api/Menu';
import { toast } from 'sonner';
import { Modal, Form, Tag, Card, Switch, Tree, Spin, Space } from 'antd';
import { handleTree } from '../utils/tree';
import type { DataNode } from 'antd/es/tree';
import { useTranslation } from 'react-i18next';

interface RoleAssignMenuFormProps {
  onSuccess?: () => void;
}

export interface RoleAssignMenuFormRef {
  open: (row: RoleVO) => void;
}

const RoleAssignMenuForm = forwardRef<RoleAssignMenuFormRef, RoleAssignMenuFormProps>(({ onSuccess }, ref) => {
  const { t } = useTranslation();
  const [dialogVisible, setDialogVisible] = useState(false);
  const [formLoading, setFormLoading] = useState(false);
  const [formData, setFormData] = useState<{
    id?: number;
    name: string;
    code: string;
    menuIds: number[];
  }>({
    id: undefined,
    name: '',
    code: '',
    menuIds: [],
  });
  const [menuOptions, setMenuOptions] = useState<DataNode[]>([]);
  const [checkedKeys, setCheckedKeys] = useState<React.Key[]>([]);
  const [halfCheckedKeys, setHalfCheckedKeys] = useState<React.Key[]>([]);
  const [expandedKeys, setExpandedKeys] = useState<React.Key[]>([]);
  const [menuExpand, setMenuExpand] = useState(false);
  const [treeNodeAll, setTreeNodeAll] = useState(false);
  const [form] = Form.useForm();

  // 获取所有节点 key（含父/子）
  const getAllKeys = (nodes: DataNode[]): React.Key[] => {
    let keys: React.Key[] = [];
    nodes.forEach((node) => {
      keys.push(node.key);
      if (node.children) {
        keys = keys.concat(getAllKeys(node.children));
      }
    });
    return keys;
  };

  // 获取所有“父节点”key（有 children 的节点）
  const getParentKeys = (nodes: DataNode[]): Set<React.Key> => {
    const parents = new Set<React.Key>();
    const walk = (ns: DataNode[]) => {
      ns.forEach((n) => {
        if (n.children && n.children.length > 0) {
          parents.add(n.key);
          walk(n.children);
        }
      });
    };
    walk(nodes);
    return parents;
  };

  // 获取某个节点的所有子节点 key（递归）
  const getChildrenKeys = (node: DataNode): React.Key[] => {
    let keys: React.Key[] = [];
    if (node.children && node.children.length > 0) {
      node.children.forEach((child) => {
        keys.push(child.key);
        keys = keys.concat(getChildrenKeys(child));
      });
    }
    return keys;
  };

  // 在树中查找节点
  const findNode = (nodes: DataNode[], key: React.Key): DataNode | null => {
    for (const node of nodes) {
      if (node.key === key) {
        return node;
      }
      if (node.children) {
        const found = findNode(node.children, key);
        if (found) return found;
      }
    }
    return null;
  };

  // 根据接口返回的ID列表，计算正确的回显状态
  // 规则：
  // 1. 如果父节点的所有子节点都在返回列表中，父节点全选，所有子节点全选
  // 2. 如果父节点只有部分子节点在返回列表中，父节点半选，只有那些子节点选中
  // 3. 如果接口返回了父节点ID，但子节点没有全部返回，不能把所有子节点都渲染为选中
  const calculateDisplayState = (returnedIds: number[], treeNodes: DataNode[]) => {
    const returnedIdSet = new Set(returnedIds);
    const checked: React.Key[] = [];
    const halfChecked: React.Key[] = [];
    const parentKeySet = getParentKeys(treeNodes);

    // 遍历返回的ID列表
    returnedIdSet.forEach((id) => {
      const node = findNode(treeNodes, id);
      if (!node) return;

      // 如果是父节点
      if (parentKeySet.has(id)) {
        const childrenKeys = getChildrenKeys(node);
        const allChildrenReturned = childrenKeys.every((childKey) => returnedIdSet.has(Number(childKey)));

        if (allChildrenReturned) {
          // 所有子节点都在返回列表中，父节点全选
          checked.push(id);
          // 所有子节点也加入 checked（如果还没加入的话）
          childrenKeys.forEach((childKey) => {
            if (!checked.includes(childKey)) {
              checked.push(childKey);
            }
          });
        } else {
          // 只有部分子节点在返回列表中，父节点半选
          halfChecked.push(id);
        }
      } else {
        // 如果是子节点，且不在 checked 中（避免重复），加入 checked
        if (!checked.includes(id)) {
          checked.push(id);
        }
      }
    });

    // 处理半选父节点的子节点：只有那些在返回列表中的子节点才应该被选中
    halfChecked.forEach((parentKey) => {
      const parentNode = findNode(treeNodes, parentKey);
      if (parentNode) {
        const childrenKeys = getChildrenKeys(parentNode);
        childrenKeys.forEach((childKey) => {
          // 只有返回列表中包含的子节点才加入 checked
          if (returnedIdSet.has(Number(childKey)) && !checked.includes(childKey)) {
            checked.push(childKey);
          }
        });
      }
    });

    return { checked, halfChecked };
  };

  // 将菜单数据转换为 Tree 组件需要的格式
  // 注意：handleTree 返回的节点类型不是 MenuVO（会带 children），这里用更宽松的类型兼容
  const convertMenuToTreeData = (menus: any[]): DataNode[] => {
    return (menus || []).map((menu: any) => ({
      title: menu?.name,
      key: menu?.id,
      children: menu?.children ? convertMenuToTreeData(menu.children) : undefined,
    }));
  };

  // 暴露方法给父组件
  useImperativeHandle(ref, () => ({
    open: async (row: RoleVO) => {
      setDialogVisible(true);
      resetForm();
      
      // 设置数据
      setFormData({
        id: row.id,
        name: row.name,
        code: row.code,
        menuIds: [],
      });

      // 加载菜单列表
      try {
        const menus = await menuApi.getSimpleMenusList();
        const menuData = (menus as any)?.data || menus;
        const treeData = handleTree(menuData || []);
        const treeNodes = convertMenuToTreeData(treeData);
        setMenuOptions(treeNodes);
        
        // 获取角色菜单权限列表
        setFormLoading(true);
        try {
          const menuIds = await roleApi.getRoleMenuList(row.id!);
          const menuIdsData = (menuIds as any)?.data || menuIds;
          setFormData(prev => ({ ...prev, menuIds: menuIdsData || [] }));
          
          // 回显：根据接口返回的ID列表，计算正确的选中状态
          // 规则：
          // 1. 如果父节点的所有子节点都在返回列表中，父节点全选，所有子节点全选
          // 2. 如果父节点只有部分子节点在返回列表中，父节点半选，只有那些子节点选中
          const returnedIds = (menuIdsData || []).map((id: number) => id);
          const { checked, halfChecked } = calculateDisplayState(returnedIds, treeNodes);
          setCheckedKeys(checked);
          setHalfCheckedKeys(halfChecked);
        } catch (error: any) {
          console.error(t('role.getRoleMenuFailed'), error);
        } finally {
          setFormLoading(false);
        }
      } catch (error: any) {
        console.error(t('role.loadMenuListFailed'), error);
        toast.error(t('role.loadMenuListFailed'));
      }
    },
  }));

  // 提交表单
  const submitForm = async () => {
    try {
      setFormLoading(true);
      
      // 获取所有选中的节点（包含半选中的父节点）
      // 说明：很多“一级菜单/父节点”在只勾选部分子节点时会处于半选中（halfChecked），
      // 若后端需要父节点ID，也必须把 halfChecked 一并提交。
      const allSelected = Array.from(new Set([...checkedKeys, ...halfCheckedKeys]));
      const menuIds = allSelected.map(key => Number(key)).filter(id => !isNaN(id));
      console.log('提交菜单权限，选中的节点IDs(含半选父节点):', menuIds);
      
      const data = {
        roleId: formData.id!,
        menuIds,
      };

      await roleApi.assignRoleMenu(data);
      toast.success(t('role.updateSuccess'));
      setDialogVisible(false);
      onSuccess?.();
    } catch (error: any) {
      toast.error(error.message || t('role.operationFailed'));
    } finally {
      setFormLoading(false);
    }
  };

  // 重置表单
  const resetForm = () => {
    setTreeNodeAll(false);
    setMenuExpand(false);
    setFormData({
      id: undefined,
      name: '',
      code: '',
      menuIds: [],
    });
    setCheckedKeys([]);
    setHalfCheckedKeys([]);
    setExpandedKeys([]);
    form.resetFields();
  };

  // 全选/全不选
  const handleCheckedTreeNodeAll = (checked: boolean) => {
    setTreeNodeAll(checked);
    if (checked) {
      // 获取所有节点的 key
      setCheckedKeys(getAllKeys(menuOptions));
      setHalfCheckedKeys([]);
    } else {
      setCheckedKeys([]);
      setHalfCheckedKeys([]);
    }
  };

  // 展开/折叠全部
  const handleCheckedTreeExpand = (expanded: boolean) => {
    setMenuExpand(expanded);
    if (expanded) {
      setExpandedKeys(getAllKeys(menuOptions));
    } else {
      setExpandedKeys([]);
    }
  };

  // 树节点选中变化
  // 注意：当 checkStrictly=false（父子联动）时，onCheck 接收的参数是数组，半选节点在 info.halfCheckedKeys
  const onCheck = (
    checkedKeysValue: React.Key[] | { checked: React.Key[]; halfChecked: React.Key[] },
    info: any
  ) => {
    // 父子联动模式下（checkStrictly=false），checkedKeysValue 是数组
    // 包含所有选中的节点（包括父节点和子节点）
    // info.halfCheckedKeys 包含半选中的父节点（部分子节点被选中）
    let checked: React.Key[] = [];
    let half: React.Key[] = [];
    
    if (Array.isArray(checkedKeysValue)) {
      // 父子联动模式
      checked = checkedKeysValue;
      half = info?.halfCheckedKeys || [];
    } else {
      // 父子不联动模式（虽然我们用的是联动模式，但类型定义需要兼容）
      checked = checkedKeysValue.checked || [];
      half = checkedKeysValue.halfChecked || [];
    }
    
    setCheckedKeys(checked);
    setHalfCheckedKeys(half);
    
    // 如果全部选中，设置全选状态
    const allKeys = getAllKeys(menuOptions);
    setTreeNodeAll(checked.length === allKeys.length && allKeys.length > 0);
  };

  return (
    <Modal
      title={t('role.menuPermission')}
      open={dialogVisible}
      onCancel={() => setDialogVisible(false)}
      onOk={submitForm}
      confirmLoading={formLoading}
      okText={t('common.confirm')}
      cancelText={t('common.cancel')}
      width={900}
      destroyOnClose
    >
      <Spin spinning={formLoading}>
        <Form
          form={form}
          layout="horizontal"
          labelCol={{ span: 3 }}
          wrapperCol={{ span: 21 }}
        >
          <Form.Item label={t('role.roleName')}>
            <Tag style={{
      fontSize: '16px',
      padding: '8px 16px', 
      lineHeight: '1',
      minHeight: '32px'
    }}>{formData.name}</Tag>
          </Form.Item>

          <Form.Item label={t('role.roleCode')}>
            <Tag style={{
      fontSize: '16px',
      padding: '8px 16px', 
      lineHeight: '1',
      minHeight: '32px'
    }}>{formData.code}</Tag>
          </Form.Item>

          <Form.Item label={t('role.menuPermission')}>
            <Card
              title={
                <Space>
                  <span>{t('role.selectAll')}</span>
                  <Switch
                    checked={treeNodeAll}
                    onChange={handleCheckedTreeNodeAll}
                    checkedChildren={t('role.yes')}
                    unCheckedChildren={t('role.no')}
                  />
                  <span style={{ marginLeft: 16 }}>{t('role.expandAll')}</span>
                  <Switch
                    checked={menuExpand}
                    onChange={handleCheckedTreeExpand}
                    checkedChildren={t('role.expand')}
                    unCheckedChildren={t('role.collapse')}
                  />
                </Space>
              }
              style={{ maxHeight: 400, overflowY: 'auto' }}
            >
              <Tree
                checkable
                checkedKeys={checkedKeys}
                expandedKeys={expandedKeys}
                onCheck={onCheck as any}
                onExpand={setExpandedKeys}
                treeData={menuOptions}
                defaultExpandAll={false}
                // 父子联动：点击父节点可以批量操作所有子节点，提高效率
                // 半选状态由 Tree 组件自动计算（根据子节点选中情况）
                checkStrictly={false}
              />
            </Card>
          </Form.Item>
        </Form>
      </Spin>
    </Modal>
  );
});

RoleAssignMenuForm.displayName = 'RoleAssignMenuForm';

export default RoleAssignMenuForm;

