package cd.casic.module.terminal.service.impl;

import cd.casic.framework.security.core.util.SecurityFrameworkUtils;
import cd.casic.module.terminal.common.Const;
import cd.casic.module.terminal.common.ErrorMessage;
import cd.casic.module.terminal.controller.dto.TerminalConnectDTO;
import cd.casic.module.terminal.dal.redis.TerminalRedisDAO;
import cd.casic.module.terminal.define.cache.TerminalCacheKeyDefine;
import cd.casic.module.terminal.host.jsch.SessionStores;
import cd.casic.module.terminal.host.terminal.dto.SftpGetContentCacheDTO;
import cd.casic.module.terminal.host.terminal.dto.SftpSetContentCacheDTO;
import cd.casic.module.terminal.host.transfer.TerminalTransferManager;
import cd.casic.module.terminal.host.transfer.session.DownloadSession;
import cd.casic.module.terminal.service.HostConnectService;
import cd.casic.module.terminal.service.TerminalSftpService;
import cn.hutool.core.io.IoUtil;
import cn.orionsec.kit.lang.constant.StandardContentType;
import cn.orionsec.kit.lang.constant.StandardHttpHeader;
import cn.orionsec.kit.lang.define.wrapper.HttpWrapper;
import cn.orionsec.kit.lang.utils.Exceptions;
import cn.orionsec.kit.lang.utils.Strings;
import cn.orionsec.kit.lang.utils.Valid;
import cn.orionsec.kit.lang.utils.io.Files1;
import cn.orionsec.kit.lang.utils.io.Streams;
import cn.orionsec.kit.net.host.SessionStore;
import cn.orionsec.kit.net.host.sftp.SftpExecutor;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Optional;

/**
 * SFTP 操作 服务实现类
 *
 * @author Jiahang Li
 * @version 1.0.0
 * @since 2024/3/4 23:35
 */
@Slf4j
@Service
public class TerminalSftpServiceImpl implements TerminalSftpService {

    @Resource
    private TerminalRedisDAO terminalRedisDAO;

    @Resource
    private HostConnectService hostConnectService;

    @Resource
    private TerminalTransferManager terminalTransferManager;

    @Override
    public void getFileContentByToken(String token, HttpServletResponse response) throws IOException {
        InputStream in = null;
        // 解析 token
        String key = TerminalCacheKeyDefine.TERMINAL_SFTP_GET_CONTENT.format(token);
        SftpGetContentCacheDTO cache = terminalRedisDAO.getJson(key, TerminalCacheKeyDefine.TERMINAL_SFTP_GET_CONTENT);
        if (cache == null) {
//            Servlets.writeHttpWrapper(response, HttpWrapper.error(ErrorMessage.FILE_ABSENT));
            HttpWrapper<Object> wrapper = HttpWrapper.error(ErrorMessage.FILE_ABSENT);
            in = Streams.toInputStream(wrapper.toJsonString().getBytes(StandardCharsets.UTF_8));
            response.setContentType(StandardContentType.APPLICATION_JSON_UTF8);
            IoUtil.copy(in, response.getOutputStream());
            return;
        }
        // 删除缓存
        terminalRedisDAO.delete(key);
        // 获取文件内容
        SessionStore sessionStore = null;
        SftpExecutor executor = null;
        try {
            // 获取终端连接信息
            TerminalConnectDTO connectInfo = hostConnectService.getSshConnectInfo(cache.getHostId(), SecurityFrameworkUtils.getLoginUser().getId());
            sessionStore = SessionStores.openSessionStore(connectInfo);
            executor = sessionStore.getSftpExecutor(connectInfo.getFileNameCharset());
            executor.connect();
            // 读取文件
            in = executor.openInputStream(cache.getPath());
            // 设置返回
//            Servlets.setContentType(response, StandardContentType.TEXT_PLAIN);
//            Servlets.transfer(response, in);
            response.setContentType(StandardContentType.TEXT_PLAIN);
            IoUtil.copy(in, response.getOutputStream());
        } catch (Exception e) {
//            Servlets.writeHttpWrapper(response, HttpWrapper.error(ErrorMessage.FILE_READ_ERROR));
            HttpWrapper<Object> wrapper = HttpWrapper.error(ErrorMessage.FILE_READ_ERROR);
            in = Streams.toInputStream(wrapper.toJsonString().getBytes(StandardCharsets.UTF_8));
            response.setContentType(StandardContentType.APPLICATION_JSON_UTF8);
            IoUtil.copy(in, response.getOutputStream());
        } finally {
            IoUtil.close(in);
            Streams.close(executor);
            Streams.close(sessionStore);
            Streams.close(in);
        }
    }

    @Override
    public void setFileContentByToken(String token, MultipartFile file) {
        // 解析 token
        String key = TerminalCacheKeyDefine.TERMINAL_SFTP_SET_CONTENT.format(token);
        SftpSetContentCacheDTO cache = terminalRedisDAO.getJson(key, TerminalCacheKeyDefine.TERMINAL_SFTP_SET_CONTENT);
        Valid.notNull(cache, ErrorMessage.FILE_ABSENT);
        // 删除缓存
        terminalRedisDAO.delete(key);
        // 写入文件内容
        SessionStore sessionStore = null;
        SftpExecutor executor = null;
        OutputStream out = null;
        InputStream in = null;
        try {
            // 获取终端连接信息
            TerminalConnectDTO connectInfo = hostConnectService.getSshConnectInfo(cache.getHostId(), SecurityFrameworkUtils.getLoginUser().getId());
            sessionStore = SessionStores.openSessionStore(connectInfo);
            executor = sessionStore.getSftpExecutor(connectInfo.getFileNameCharset());
            executor.connect();
            // 写入文件
            out = executor.openOutputStream(cache.getPath());
            Streams.transfer(in = file.getInputStream(), out);
        } catch (Exception e) {
            throw Exceptions.app(ErrorMessage.OPERATE_ERROR);
        } finally {
            Streams.close(executor);
            Streams.close(sessionStore);
            Streams.close(out);
            Streams.close(in);
        }
    }

    @Override
    public StreamingResponseBody downloadWithTransferToken(String channelId, String transferToken, HttpServletResponse response) {
        // 获取会话
        DownloadSession session = (DownloadSession) Optional.ofNullable(channelId)
                .map(terminalTransferManager::getHandler)
                .map(s -> s.getSessionByToken(transferToken))
                .filter(s -> s instanceof DownloadSession)
                .orElse(null);
        // 响应会话
        if (session == null) {
            response.setContentType(StandardContentType.TEXT_HTML);
            response.setCharacterEncoding(Const.UTF_8);
            return outputStream -> outputStream.write(Strings.bytes(ErrorMessage.SESSION_ABSENT));
        }
        // 响应文件
        setAttachmentHeader(response, Files1.getFileName(session.getPath()));
        return session;
    }

    public void setAttachmentHeader(HttpServletResponse response, String fileName) {
        response.setContentType(StandardContentType.APPLICATION_STREAM);
        response.setHeader(StandardHttpHeader.ACCESS_CONTROL_EXPOSE_HEADERS, StandardHttpHeader.CONTENT_DISPOSITION);
        try {
            response.setHeader(StandardHttpHeader.CONTENT_DISPOSITION, "attachment; filename=" + URLEncoder.encode(fileName, cn.orionsec.kit.lang.constant.Const.UTF_8));
        } catch (UnsupportedEncodingException e) {
            throw Exceptions.unsupportedEncoding(e);
        }
    }
}
