#include "camera_pipeline.hpp"
#include "bounded_queue.hpp"
#include "latest_result.hpp"
#include "detector.hpp"
#include "door_risk.hpp"
#include "records.hpp"
#include "inference_session.hpp"
#include "can_bus.hpp"
#include "serialization.hpp"
#include "result_conversion.hpp"
#include <opencv2/core.hpp>
#include <atomic>
#include <cmath>
#include <cstring>
#include <exception>
#include <mutex>
#include <limits>
#include <thread>

namespace cms {
namespace {
struct Stream {
    cudaStream_t value{};
    Stream(){CUDA_CHECK(cudaStreamCreateWithFlags(&value,cudaStreamNonBlocking));}
    ~Stream(){cudaStreamDestroy(value);}
};
cv::Mat imageView(BgrImage const& image) {
    return cv::Mat(image.height,image.width,CV_8UC3,const_cast<std::uint8_t*>(image.pixels.data()),image.stride);
}
double elapsedMs(std::uint64_t start){return (monotonicNs()-start)*1e-6;}
void attachPoses(FrameResult& result,fangjia::PoseFrameResult const& fused) {
    for(std::size_t i=0;i<fused.selectedPeople.size();++i) {
        auto const& source=fused.selectedPeople[i];
        PersonResult person;person.track_id=source.box.trackId;person.confidence=source.box.score;
        person.box={source.box.x1,source.box.y1,source.box.x2,source.box.y2};
        for(std::size_t k=0;k<133;++k) {
            auto const&a=source.keypoints[k];auto const&b=fused.smoothedPeople[i].keypoints[k];auto const&c=fused.v2People[i].keypoints[k];
            person.selected[k]={a.x,a.y,a.score};person.smoothed[k]={b.x,b.y,b.score};person.temporal[k]={c.x,c.y,c.score};
        }
        for(int side=0;side<2;++side){auto const& m=fused.metadata[i][side];person.hands[side]={m.handState,m.handSource,m.limbConsensus,m.distanceEligible,m.alarmEligible};}
        result.persons.push_back(std::move(person));
    }
    for(auto const& source:fused.globalHands) {
        HandRegion hand;hand.person_id=source.personId;hand.track_id=source.trackId;hand.side=source.left?"left":"right";
        for(std::size_t k=0;k<21;++k)hand.keypoints[k]={source.keypoints[k].x,source.keypoints[k].y,source.keypoints[k].score};
        result.hand_regions.push_back(std::move(hand));
    }
}
}
struct CameraPipeline::Impl {
    struct Detected {CapturedFrame frame;std::vector<fangjia::Box> boxes;double elapsed_ms{};};
    PipelineConfig config;
    BoundedQueue<CapturedFrame> captures;
    BoundedQueue<Detected> detections;
    LatestResult results;
    std::unique_ptr<OutputWriter> outputWriter;
    std::unique_ptr<SessionWriter> sessionWriter;
    std::unique_ptr<SessionReader> sessionReader;
    fangjia::TrtLogger logger;
    std::unique_ptr<Detector> detector;
    std::unique_ptr<fangjia::RtmwEngine> pose;
    std::unique_ptr<DoorRisk> risk;
    std::shared_ptr<const BinaryMask> doorPixels;
    std::unique_ptr<Camera> camera;
    std::unique_ptr<CanBus> can;
    std::thread captureThread,detectorThread,inferenceThread;
    mutable std::mutex stateMutex;
    std::exception_ptr failure;
    bool started{};
    std::atomic<bool> stopped{false},active{false};
    std::atomic<std::uint64_t> captured{0},published{0},gaps{0},errorFrames{0};
    double totalLatency{};

    explicit Impl(PipelineConfig c):config(std::move(c)),captures(config.queue_capacity),detections(config.queue_capacity) {
        if(!config.replay_session.empty())sessionReader=std::make_unique<SessionReader>(config);
        auto canConfig=config.can;if(sessionReader)canConfig.enabled=false;
        can=std::make_unique<CanBus>(canConfig);
        if(config.queue_capacity>8)throw std::invalid_argument("Queue capacity must be between 1 and 8");
        if(config.max_people<1||config.max_people>16||!std::isfinite(config.detector_confidence)||config.detector_confidence<=0||config.detector_confidence>=1)
            throw std::invalid_argument("Invalid detector configuration");
        if(!std::filesystem::is_regular_file(config.door_mask))throw std::invalid_argument("Missing door mask: "+config.door_mask.string());
    }
    ~Impl(){stop();joinWorkers();}
    void joinWorkers() noexcept {
        if(captureThread.joinable())captureThread.join();
        if(detectorThread.joinable())detectorThread.join();
        if(inferenceThread.joinable())inferenceThread.join();
        can->join();
    }
    void stop() noexcept {
        stopped=true;active=false;
        can->requestStop();
        // Camera stays alive until all threads have joined. start()/destruction
        // are owner-thread operations; requestStop() only interrupts existing work.
        {std::lock_guard<std::mutex> lock(stateMutex);if(camera)camera->requestStop();}
        captures.close(true);detections.close(true);
    }
    void fail() noexcept {
        {std::lock_guard<std::mutex> lock(stateMutex);if(!failure)failure=std::current_exception();}
        stop();
    }
    void rethrowFailure() const {
        std::lock_guard<std::mutex> lock(stateMutex);if(failure)std::rethrow_exception(failure);
    }
    void start() {
        {std::lock_guard<std::mutex> lock(stateMutex);if(started||stopped)throw std::logic_error("CameraPipeline can only be started once");started=true;}
        try {
            can->start();
            outputWriter=std::make_unique<OutputWriter>(config.output);
            sessionWriter=std::make_unique<SessionWriter>(config);
            detector=std::make_unique<Detector>(config.engines/"yolo11s.engine",logger,config.detector_confidence,config.max_people);
            pose=std::make_unique<fangjia::RtmwEngine>(config.engines/"rtmw_l_384.engine",logger,config.camera.width,config.camera.height,config.max_people);
            DoorOptions options;options.constant=config.door_mask;options.warning=config.warning_relative;options.release=config.release_relative;
            risk=std::make_unique<DoorRisk>(options,cv::Size(config.camera.width,config.camera.height));risk->initialize();
            if(risk->mask().empty()||cv::countNonZero(risk->mask())==0)throw std::runtime_error("Door mask has no usable region");
            auto pixels=std::make_shared<BinaryMask>();pixels->width=config.camera.width;pixels->height=config.camera.height;pixels->stride=config.camera.width;
            auto const& mask=risk->mask();pixels->pixels.resize(std::size_t(mask.cols)*mask.rows);
            for(int y=0;y<mask.rows;++y)std::memcpy(pixels->pixels.data()+std::size_t(y)*mask.cols,mask.ptr(y),mask.cols);
            doorPixels=std::move(pixels);
            Stream warmup;
            pose->warmup(warmup.value);
            detector->infer(cv::Mat(config.camera.height,config.camera.width,CV_8UC3,cv::Scalar()),warmup.value);
            {std::lock_guard<std::mutex> lock(stateMutex);if(!sessionReader)camera=std::make_unique<Camera>(config.camera);if(stopped&&camera)camera->requestStop();}
            active=!stopped.load();
            captureThread=std::thread([this]{try{captureLoop();}catch(...){fail();}});
            detectorThread=std::thread([this]{try{detectorLoop();}catch(...){fail();}});
            inferenceThread=std::thread([this]{try{inferenceLoop();}catch(...){fail();}});
        }catch(...){stop();joinWorkers();throw;}
    }
    void captureLoop() {
        while(!stopped) {
            auto frame=sessionReader?sessionReader->read():camera->read();if(!frame)break;
            gaps+=frame->metadata.skipped_driver_frames;
            errorFrames+=frame->metadata.discarded_error_frames;
            if(!captures.push(std::move(*frame)))break;
            ++captured;
        }
        captures.close();
    }
    void detectorLoop() {
        Stream stream;
        while(!stopped) {
            auto frame=captures.pop();if(!frame)break;
            auto startedAt=monotonicNs();auto boxes=detector->infer(imageView(*frame->image),stream.value);
            if(!detections.push({std::move(*frame),std::move(boxes),elapsedMs(startedAt)}))break;
        }
        detections.close();
    }
    void inferenceLoop() {
        Stream stream;fangjia::BoxTracker boxesTracker;fangjia::PoseFusionTracker fusion;
        while(!stopped) {
            auto detected=detections.pop();if(!detected)break;
            if(detected->frame.metadata.frame_index>static_cast<std::uint64_t>(std::numeric_limits<int>::max()))throw std::runtime_error("Frame index limit reached; restart pipeline");
            auto index=static_cast<int>(detected->frame.metadata.frame_index);
            auto& boxes=detected->boxes;
            auto frame=imageView(*detected->frame.image);
            boxesTracker.update(boxes,index,frame.cols,frame.rows);
            auto poseStart=monotonicNs();
            auto raw=pose->infer(frame,boxes,stream.value);auto isolated=raw;
            for(std::size_t p=0;p<boxes.size();++p) {
                auto const& box=boxes[p];double area=std::max(double(box.x2-box.x1)*(box.y2-box.y1),1.0);
                bool needsIsolation=false;
                for(std::size_t q=0;q<boxes.size();++q)if(q!=p&&std::max(double(boxes[q].x2-boxes[q].x1)*(boxes[q].y2-boxes[q].y1),1.0)<0.5*area)needsIsolation=true;
                if(needsIsolation){auto masked=maskOtherPeople(frame,boxes,p);isolated[p]=pose->infer(masked,{box},stream.value).at(0);}
            }
            auto poseElapsed=elapsedMs(poseStart);auto fusionStart=monotonicNs();
            auto fused=fusion.update(raw,isolated,index);
            auto record=makeRecord(index,raw,isolated,fused,!config.record_session.empty(),false);
            auto riskData=risk->evaluate(index,record);
            auto result=std::make_shared<FrameResult>(makeFrameResult(record,riskData,detected->frame.metadata));
            if(detected->frame.replay_record) {
                auto const& recorded=detected->frame.replay_record->at("result");
                if(recorded.contains("can"))result->can=deserializeCan(recorded.at("can"));
                else result->can.status="NOT_RECORDED";
                result->can.source="replay";
            }else {
                if(!sessionReader&&config.can.enabled)can->queueRisk(result->risk.risk_level);
                result->can=can->snapshot();
            }
            attachPoses(*result,fused);
            result->image=std::move(detected->frame.image);result->door_mask_pixels=doorPixels;
            result->door_mask_calibration_verified=config.door_mask_calibration_verified;
            result->timing={detected->elapsed_ms,poseElapsed,elapsedMs(fusionStart),elapsedMs(result->capture.capture_monotonic_ns)};
            FrameResultPtr immutable=result;
            results.publish(std::move(immutable));
            {std::lock_guard<std::mutex> lock(stateMutex);totalLatency+=result->timing.capture_to_result_ms;++published;}
            if(!config.record_session.empty())sessionWriter->write(*result,record,detected->frame.replay_record.get());
            outputWriter->write(*result);
        }
        sessionWriter->finish();
        active=false;
    }
};
CameraPipeline::CameraPipeline(PipelineConfig config):impl_(std::make_unique<Impl>(std::move(config))){}
CameraPipeline::~CameraPipeline()=default;
void CameraPipeline::start(){impl_->start();}
FrameResultPtr CameraPipeline::latest()const{return impl_->results.load();}
CanSnapshot CameraPipeline::canStatus()const {
    if(impl_->sessionReader) {
        auto result=impl_->results.load();if(result)return result->can;
        CanSnapshot s;s.source="replay";s.status="WAITING";return s;
    }
    return impl_->can->snapshot();
}
bool CameraPipeline::sendCan(CanFrame const& frame) {
    if(impl_->sessionReader)throw std::logic_error("CAN transmission is disabled during replay");
    return impl_->can->send(frame);
}
bool CameraPipeline::running()const noexcept{return impl_->active.load();}
void CameraPipeline::requestStop()noexcept{impl_->stop();}
void CameraPipeline::join(){impl_->joinWorkers();{std::lock_guard<std::mutex> lock(impl_->stateMutex);impl_->camera.reset();}impl_->rethrowFailure();}
PipelineStats CameraPipeline::stats()const {
    std::lock_guard<std::mutex> lock(impl_->stateMutex);auto count=impl_->published.load();
    return {impl_->captured.load(),count,impl_->gaps.load(),impl_->errorFrames.load(),count?impl_->totalLatency/count:0};
}
}