beeta/backend/src/controllers/StreamController.cpp

#include "StreamController.h"
#include "../services/DatabaseService.h"
#include "../services/StatsService.h"
#include "../services/RedisHelper.h"
#include "../services/OmeClient.h"
#include "../services/AuthService.h"
#include <drogon/utils/Utilities.h>
#include <drogon/Cookie.h>
#include <random>
#include <sstream>
#include <iomanip>
#include <chrono>

using namespace drogon::orm;

// Helper functions at the top
namespace {
    // JSON response helper - saves 6-8 lines per endpoint
    HttpResponsePtr jsonResp(const Json::Value& j, HttpStatusCode c = k200OK) {
        auto r = HttpResponse::newHttpJsonResponse(j);
        r->setStatusCode(c);
        return r;
    }
    
    HttpResponsePtr jsonOk(const Json::Value& data) {
        return jsonResp(data);
    }
    
    HttpResponsePtr jsonError(const std::string& error, HttpStatusCode code = k400BadRequest) {
        Json::Value j;
        j["success"] = false;
        j["error"] = error;
        return jsonResp(j, code);
    }
    
    // Quick JSON builder for common patterns
    Json::Value json(std::initializer_list<std::pair<const char*, Json::Value>> items) {
        Json::Value j;
        for (const auto& [key, value] : items) {
            j[key] = value;
        }
        return j;
    }
    
    UserInfo getUserFromRequest(const HttpRequestPtr &req) {
        UserInfo user;
        std::string auth = req->getHeader("Authorization");
        
        if (auth.empty() || auth.substr(0, 7) != "Bearer ") {
            return user;
        }
        
        std::string token = auth.substr(7);
        AuthService::getInstance().validateToken(token, user);
        return user;
    }
}

// Static member definitions
std::mutex StreamWebSocketController::connectionsMutex_;
std::unordered_map<std::string, std::unordered_set<WebSocketConnectionPtr>> StreamWebSocketController::tokenConnections_;
std::unordered_set<WebSocketConnectionPtr> StreamWebSocketController::connections_;

void StreamController::health(const HttpRequestPtr &,
                             std::function<void(const HttpResponsePtr &)> &&callback) {
    callback(jsonOk(json({
        {"status", "ok"},
        {"timestamp", Json::Int64(std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now().time_since_epoch()
        ).count())}
    })));
}

void StreamController::validateStreamKey(const HttpRequestPtr &,
                                        std::function<void(const HttpResponsePtr &)> &&callback,
                                        const std::string &key) {
    // Now validate against realms table
    auto dbClient = app().getDbClient();
    *dbClient << "SELECT 1 FROM realms WHERE stream_key = $1 AND is_active = true"
              << key
              >> [callback, key](const Result &r) {
                  bool valid = !r.empty();
                  if (valid) {
                      // Store in Redis
                      RedisHelper::storeKey("stream_key:" + key, "1", 86400);
                  }
                  callback(jsonOk(json({{"valid", valid}})));
              }
              >> [callback](const DrogonDbException &e) {
                  LOG_ERROR << "Database error: " << e.base().what();
                  callback(jsonOk(json({{"valid", false}})));
              };
}

void StreamController::disconnectStream(const HttpRequestPtr &req,
                                       std::function<void(const HttpResponsePtr &)> &&callback,
                                       const std::string &streamKey) {
    UserInfo user = getUserFromRequest(req);
    if (user.id == 0) {
        callback(jsonError("Unauthorized", k401Unauthorized));
        return;
    }
    
    // Check if user owns this stream or is admin
    auto dbClient = app().getDbClient();
    *dbClient << "SELECT user_id FROM realms WHERE stream_key = $1 AND is_active = true"
              << streamKey
              >> [user, callback, streamKey](const Result& r) {
                  if (r.empty() || (r[0]["user_id"].as<int64_t>() != user.id && !user.isAdmin)) {
                      callback(jsonError("Forbidden", k403Forbidden));
                      return;
                  }
                  
                  OmeClient::getInstance().disconnectStream(streamKey, [callback](bool success) {
                      if (success) {
                          callback(jsonOk(json({
                              {"success", true},
                              {"message", "Stream disconnected"}
                          })));
                      } else {
                          callback(jsonError("Failed to disconnect stream"));
                      }
                  });
              }
              >> [callback](const DrogonDbException& e) {
                  LOG_ERROR << "Database error: " << e.base().what();
                  callback(jsonError("Database error"));
              };
}

void StreamController::getStreamStats(const HttpRequestPtr &,
                                     std::function<void(const HttpResponsePtr &)> &&callback,
                                     const std::string &streamKey) {
    StatsService::getInstance().getStreamStats(streamKey,
        [callback](bool success, const StreamStats& stats) {
            if (success) {
                Json::Value json;
                json["success"] = true;
                
                auto& s = json["stats"];
                s["connections"] = static_cast<Json::Int64>(stats.uniqueViewers);
                s["total_connections"] = static_cast<Json::Int64>(stats.totalConnections);
                s["bytes_in"] = static_cast<Json::Int64>(stats.totalBytesIn);
                s["bytes_out"] = static_cast<Json::Int64>(stats.totalBytesOut);
                s["bitrate"] = stats.bitrate;
                s["codec"] = stats.codec;
                s["resolution"] = stats.resolution;
                s["fps"] = stats.fps;
                s["is_live"] = stats.isLive;
                
                if (stats.totalBytesIn > 0) {
                    s["data_rate_in"] = stats.bitrate / 1000.0;
                }
                if (stats.totalBytesOut > 0) {
                    s["data_rate_out"] = stats.totalBytesOut / 1024.0 / 1024.0;
                }
                
                // Protocol breakdown
                auto& pc = s["protocol_connections"];
                pc["webrtc"] = static_cast<Json::Int64>(stats.protocolConnections.webrtc);
                pc["hls"] = static_cast<Json::Int64>(stats.protocolConnections.hls);
                pc["llhls"] = static_cast<Json::Int64>(stats.protocolConnections.llhls);
                pc["dash"] = static_cast<Json::Int64>(stats.protocolConnections.dash);
                
                callback(jsonResp(json));
            } else {
                callback(jsonError("Failed to retrieve stream stats"));
            }
        });
}

void StreamController::getActiveStreams(const HttpRequestPtr &,
                                       std::function<void(const HttpResponsePtr &)> &&callback) {
    OmeClient::getInstance().getActiveStreams([callback](bool success, const Json::Value& omeResponse) {
        if (success) {
            LOG_INFO << "Active streams: " << omeResponse["response"].toStyledString();
            callback(jsonOk(json({
                {"success", true},
                {"streams", omeResponse["response"]}
            })));
        } else {
            callback(jsonError("Failed to get active streams from OME"));
        }
    });
}

void StreamController::issueViewerToken(const HttpRequestPtr &,
                                       std::function<void(const HttpResponsePtr &)> &&callback,
                                       const std::string &streamKey) {
    // Validate against realms
    auto dbClient = app().getDbClient();
    *dbClient << "SELECT 1 FROM realms WHERE stream_key = $1 AND is_active = true"
              << streamKey
              >> [callback, streamKey](const Result& r) {
                  if (r.empty()) {
                      callback(jsonResp({}, k404NotFound));
                      return;
                  }
                  
                  auto bytes = drogon::utils::genRandomString(32);
                  std::string token = drogon::utils::base64Encode(
                      (const unsigned char*)bytes.data(), bytes.length()
                  );
                  
                  RedisHelper::storeKeyAsync("viewer_token:" + token, streamKey, 30,
                      [callback, token](bool stored) {
                          if (!stored) {
                              callback(jsonResp({}, k500InternalServerError));
                              return;
                          }
                          
                          auto resp = HttpResponse::newHttpResponse();
                          
                          Cookie cookie("viewer_token", token);
                          cookie.setPath("/");
                          cookie.setHttpOnly(true);
                          cookie.setSecure(false);
                          cookie.setMaxAge(300);
                          resp->addCookie(cookie);
                          
                          Json::Value body;
                          body["success"] = true;
                          body["viewer_token"] = token;
                          body["expires_in"] = 30;
                          
                          resp->setContentTypeCode(CT_APPLICATION_JSON);
                          resp->setBody(Json::FastWriter().write(body));
                          
                          callback(resp);
                      }
                  );
              }
              >> [callback](const DrogonDbException& e) {
                  LOG_ERROR << "Database error: " << e.base().what();
                  callback(jsonResp({}, k500InternalServerError));
              };
}

void StreamController::heartbeat(const HttpRequestPtr &req,
                                std::function<void(const HttpResponsePtr &)> &&callback,
                                const std::string &streamKey) {
    auto token = req->getCookie("viewer_token");
    if (token.empty()) {
        callback(jsonResp({}, k403Forbidden));
        return;
    }
    
    RedisHelper::getKeyAsync("viewer_token:" + token, 
        [callback, streamKey, token](const std::string& storedStreamKey) {
            if (storedStreamKey != streamKey) {
                callback(jsonResp({}, k403Forbidden));
                return;
            }
            
            services::RedisHelper::instance().expireAsync("viewer_token:" + token, 30,
                [callback](bool success) {
                    if (!success) {
                        callback(jsonResp({}, k500InternalServerError));
                        return;
                    }
                    
                    callback(jsonOk(json({
                        {"success", true},
                        {"renewed", true}
                    })));
                }
            );
        }
    );
}

// WebSocket implementation
void StreamWebSocketController::handleNewMessage(const WebSocketConnectionPtr&,
                                                std::string &&message,
                                                const WebSocketMessageType &type) {
    if (type == WebSocketMessageType::Text) {
        Json::Value msg;
        Json::Reader reader;
        if (reader.parse(message, msg) && msg["type"].asString() == "subscribe") {
            LOG_INFO << "Client subscribed to stream updates";
        }
    }
}

void StreamWebSocketController::handleNewConnection(const HttpRequestPtr &req,
                                                   const WebSocketConnectionPtr& wsConnPtr) {
    LOG_INFO << "New WebSocket connection established";
    
    auto token = req->getCookie("viewer_token");
    if (token.empty()) {
        LOG_WARN << "WebSocket connection without viewer token";
        wsConnPtr->shutdown();
        return;
    }
    
    RedisHelper::getKeyAsync("viewer_token:" + token,
        [wsConnPtr, token](const std::string& streamKey) {
            if (streamKey.empty()) {
                LOG_WARN << "Invalid viewer token";
                wsConnPtr->shutdown();
                return;
            }
            
            std::lock_guard<std::mutex> lock(connectionsMutex_);
            tokenConnections_[token].insert(wsConnPtr);
            connections_.insert(wsConnPtr);
            
            LOG_INFO << "WebSocket authenticated for stream: " << streamKey;
        }
    );
}

void StreamWebSocketController::handleConnectionClosed(const WebSocketConnectionPtr& wsConnPtr) {
    LOG_INFO << "WebSocket connection closed";
    
    std::lock_guard<std::mutex> lock(connectionsMutex_);
    
    std::string tokenToDelete;
    for (auto& [token, conns] : tokenConnections_) {
        if (conns.erase(wsConnPtr)) {
            if (conns.empty()) {
                tokenToDelete = token;
            }
            break;
        }
    }
    
    connections_.erase(wsConnPtr);
    
    if (!tokenToDelete.empty()) {
        tokenConnections_.erase(tokenToDelete);
        
        RedisHelper::deleteKeyAsync("viewer_token:" + tokenToDelete, 
            [tokenToDelete](bool success) {
                if (success) {
                    LOG_INFO << "Deleted viewer token on disconnect: " << tokenToDelete;
                } else {
                    LOG_WARN << "Failed to delete viewer token: " << tokenToDelete;
                }
            }
        );
    }
}

void StreamWebSocketController::broadcastKeyUpdate(const std::string& userId, const std::string& newKey) {
    Json::Value msg;
    msg["type"] = "key_regenerated";
    msg["user_id"] = userId;
    msg["stream_key"] = newKey;
    
    auto msgStr = Json::FastWriter().write(msg);
    
    std::lock_guard<std::mutex> lock(connectionsMutex_);
    for (const auto& conn : connections_) {
        if (conn->connected()) {
            conn->send(msgStr);
        }
    }
}

void StreamWebSocketController::broadcastStatsUpdate(const Json::Value& stats) {
    std::string jsonStr = Json::FastWriter().write(stats);
    
    std::lock_guard<std::mutex> lock(connectionsMutex_);
    for (const auto& conn : connections_) {
        if (conn->connected()) {
            conn->send(jsonStr);
        }
    }
}