#include "file_transfer.h"

#include <atomic>
#include <chrono>

#include "rd_log.h"

namespace crossdesk {

namespace {
std::atomic<uint32_t> g_next_file_id{1};
}  // namespace

uint32_t FileSender::NextFileId() { return g_next_file_id.fetch_add(1); }

int FileSender::SendFile(const std::filesystem::path& path,
                         const std::string& label, const SendFunc& send,
                         std::size_t chunk_size, uint32_t file_id) {
  if (!send) {
    LOG_ERROR("FileSender::SendFile: send function is empty");
    return -1;
  }

  std::error_code ec;
  if (!std::filesystem::exists(path, ec) ||
      !std::filesystem::is_regular_file(path, ec)) {
    LOG_ERROR("FileSender::SendFile: file [{}] not found or not regular",
              path.string().c_str());
    return -1;
  }

  uint64_t total_size = std::filesystem::file_size(path, ec);
  if (ec) {
    LOG_ERROR("FileSender::SendFile: failed to get size of [{}]: {}",
              path.string().c_str(), ec.message().c_str());
    return -1;
  }

  std::ifstream ifs(path, std::ios::binary);
  if (!ifs.is_open()) {
    LOG_ERROR("FileSender::SendFile: failed to open [{}]",
              path.string().c_str());
    return -1;
  }
  LOG_INFO("FileSender send file {}, total size {}, file_id={}",
           path.string().c_str(), total_size, file_id);

  if (file_id == 0) {
    file_id = NextFileId();
  }
  const uint32_t final_file_id = file_id;
  uint64_t offset = 0;
  bool is_first = true;
  std::string file_name = label.empty() ? path.filename().string() : label;

  std::vector<char> buffer;
  buffer.resize(chunk_size);

  while (ifs && offset < total_size) {
    uint64_t remaining = total_size - offset;
    uint32_t to_read =
        static_cast<uint32_t>(std::min<uint64_t>(remaining, chunk_size));

    ifs.read(buffer.data(), static_cast<std::streamsize>(to_read));
    std::streamsize bytes_read = ifs.gcount();
    if (bytes_read <= 0) {
      break;
    }

    bool is_last = (offset + static_cast<uint64_t>(bytes_read) >= total_size);
    const std::string* name_ptr = is_first ? &file_name : nullptr;

    std::vector<char> chunk = BuildChunk(
        final_file_id, offset, total_size, buffer.data(),
        static_cast<uint32_t>(bytes_read), name_ptr, is_first, is_last);

    int ret = send(chunk.data(), chunk.size());
    if (ret != 0) {
      LOG_ERROR("FileSender::SendFile: send failed for [{}], ret={}",
                path.string().c_str(), ret);
      return ret;
    }

    offset += static_cast<uint64_t>(bytes_read);
    is_first = false;
  }

  return 0;
}

std::vector<char> FileSender::BuildChunk(uint32_t file_id, uint64_t offset,
                                         uint64_t total_size, const char* data,
                                         uint32_t data_size,
                                         const std::string* file_name,
                                         bool is_first, bool is_last) {
  return protocol::EncodeFileChunk(file_id, offset, total_size, data,
                                   data_size, file_name, is_first, is_last);
}

// ---------- FileReceiver ----------

FileReceiver::FileReceiver() : output_dir_(GetDefaultDesktopPath()) {}

FileReceiver::FileReceiver(const std::filesystem::path& output_dir)
    : output_dir_(output_dir) {
  std::error_code ec;
  if (!output_dir_.empty()) {
    std::filesystem::create_directories(output_dir_, ec);
    if (ec) {
      LOG_ERROR("FileReceiver: failed to create output dir [{}]: {}",
                output_dir_.string().c_str(), ec.message().c_str());
    }
  }
}

std::filesystem::path FileReceiver::GetDefaultDesktopPath() {
#ifdef _WIN32
  const char* home_env = std::getenv("USERPROFILE");
#else
  const char* home_env = std::getenv("HOME");
#endif
  if (!home_env) {
    return std::filesystem::path{};
  }

  std::filesystem::path desktop_path =
      std::filesystem::path(home_env) / "Desktop";

  std::error_code ec;
  std::filesystem::create_directories(desktop_path, ec);
  if (ec) {
    LOG_ERROR("FileReceiver: failed to create desktop directory [{}]: {}",
              desktop_path.string().c_str(), ec.message().c_str());
  }

  return desktop_path;
}

bool FileReceiver::OnData(const char* data, size_t size) {
  protocol::FileChunkView chunk;
  if (!protocol::DecodeFileChunk(data, size, &chunk)) {
    LOG_ERROR("FileReceiver::OnData: invalid buffer");
    return false;
  }
  const std::string* file_name =
      chunk.file_name.empty() ? nullptr : &chunk.file_name;
  return HandleChunk(chunk.header, chunk.payload, chunk.payload_size,
                     file_name);
}

bool FileReceiver::HandleChunk(const FileChunkHeader& header,
                               const char* payload, size_t payload_size,
                               const std::string* file_name) {
  auto it = contexts_.find(header.file_id);
  if (it == contexts_.end()) {
    // new file context must start with first chunk.
    if ((header.flags & 0x01) == 0) {
      LOG_ERROR("FileReceiver: received non-first chunk for unknown file_id={}",
                header.file_id);
      return false;
    }

    FileContext ctx;
    ctx.total_size = header.total_size;

    std::string filename;
    if (file_name && !file_name->empty()) {
      filename = *file_name;
    } else {
      filename = "received_" + std::to_string(header.file_id);
    }

    ctx.file_name = filename;

    std::filesystem::path save_path = output_dir_.empty()
                                          ? std::filesystem::path(filename)
                                          : output_dir_ / filename;

    // if file exists, append timestamp.
    std::error_code ec;
    if (std::filesystem::exists(save_path, ec)) {
      auto now = std::chrono::system_clock::now();
      auto ts = std::chrono::duration_cast<std::chrono::milliseconds>(
                    now.time_since_epoch())
                    .count();
      save_path = save_path.parent_path() /
                  (save_path.stem().string() + "_" + std::to_string(ts) +
                   save_path.extension().string());
    }

    ctx.ofs.open(save_path, std::ios::binary | std::ios::trunc);
    if (!ctx.ofs.is_open()) {
      LOG_ERROR("FileReceiver: failed to open [{}] for writing",
                save_path.string().c_str());
      return false;
    }

    contexts_.emplace(header.file_id, std::move(ctx));
    it = contexts_.find(header.file_id);
  }

  FileContext& ctx = it->second;

  if (payload_size > 0 && payload) {
    ctx.ofs.seekp(static_cast<std::streamoff>(header.offset), std::ios::beg);
    ctx.ofs.write(payload, static_cast<std::streamsize>(payload_size));
    if (!ctx.ofs.good()) {
      LOG_ERROR("FileReceiver: write failed for file_id={}", header.file_id);
      return false;
    }
    ctx.received += static_cast<uint64_t>(payload_size);
  }

  // Send ACK after processing chunk
  if (on_send_ack_) {
    FileTransferAck ack{};
    ack.magic = kFileAckMagic;
    ack.file_id = header.file_id;
    ack.acked_offset = header.offset + static_cast<uint64_t>(payload_size);
    ack.total_size = header.total_size;
    ack.flags = 0;

    bool is_last = (header.flags & 0x02) != 0;
    if (is_last || ctx.received >= ctx.total_size) {
      ack.flags |= 0x01;  // completed
    }

    int ret = on_send_ack_(ack);
    if (ret != 0) {
      LOG_ERROR("FileReceiver: failed to send ACK for file_id={}, ret={}",
                header.file_id, ret);
    }
  }

  bool is_last = (header.flags & 0x02) != 0;
  if (is_last || ctx.received >= ctx.total_size) {
    ctx.ofs.close();

    std::filesystem::path saved_path =
        output_dir_.empty() ? std::filesystem::path(ctx.file_name)
                            : output_dir_ / ctx.file_name;

    LOG_INFO("FileReceiver: file received complete, file_id={}, size={}",
             header.file_id, ctx.received);

    contexts_.erase(header.file_id);
  }

  return true;
}

}  // namespace crossdesk