#include "visualizer.hpp"
#include "camera_pipeline.hpp"
#include <opencv2/highgui.hpp>
#include <opencv2/highgui/highgui_c.h>
#include <opencv2/imgproc.hpp>
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <sstream>
#include <stdexcept>

namespace cms {
namespace {
constexpr char windowName[]="CMS Camera Risk";
cv::Point point(PosePoint const& p){return {int(std::nearbyint(p.x)),int(std::nearbyint(p.y))};}
cv::Scalar color(std::string const& level) {
    if(level=="danger")return {40,40,255};
    if(level=="warning")return {0,160,255};
    if(level=="watch")return {0,240,240};
    if(level=="clear")return {80,220,90};
    return {190,190,190};
}
void segment(cv::Mat& image,cv::Point a,cv::Point b,cv::Scalar c,int thickness,bool dashed) {
    if(!dashed){cv::line(image,a,b,c,thickness,cv::LINE_AA);return;}
    cv::Point2d d=b-a;double length=cv::norm(d);if(length<1)return;d*=1.0/length;
    for(double x=0;x<length;x+=14)cv::line(image,cv::Point2d(a)+d*x,cv::Point2d(a)+d*std::min(x+8,length),c,thickness,cv::LINE_AA);
}
}
cv::Mat renderResult(FrameResult const& r) {
    if(!r.image)throw std::invalid_argument("Result has no frame image");
    auto const& image=*r.image;
    if(image.width<=0||image.height<=0||image.stride<image.width*3||image.pixels.size()<std::size_t(image.stride)*image.height)
        throw std::invalid_argument("Malformed result image");
    cv::Mat view(image.height,image.width,CV_8UC3,const_cast<std::uint8_t*>(image.pixels.data()),image.stride);
    auto canvas=view.clone();
    if(r.door_mask_pixels) {
        auto const& m=*r.door_mask_pixels;
        if(m.width!=image.width||m.height!=image.height||m.stride<m.width||m.pixels.size()<std::size_t(m.stride)*m.height)
            throw std::invalid_argument("Result mask does not match image");
        cv::Mat mask(m.height,m.width,CV_8U,const_cast<std::uint8_t*>(m.pixels.data()),m.stride);
        auto tint=canvas.clone();tint.setTo(cv::Scalar(255,180,0),mask);
        cv::addWeighted(canvas,0.75,tint,0.25,0,canvas);
        std::vector<std::vector<cv::Point>> contours;cv::findContours(mask.clone(),contours,cv::RETR_EXTERNAL,cv::CHAIN_APPROX_SIMPLE);
        cv::drawContours(canvas,contours,-1,{255,180,0},2);
    }
    for(auto const& person:r.persons) {
        auto const& b=person.box;
        cv::rectangle(canvas,{int(b[0]),int(b[1])},{int(b[2]),int(b[3])},{0,220,255},2);
        std::ostringstream text;text<<"ID "<<person.track_id<<" "<<std::fixed<<std::setprecision(2)<<person.confidence;
        cv::putText(canvas,text.str(),{int(b[0])+4,std::max(180,int(b[1])+25)},cv::FONT_HERSHEY_SIMPLEX,.8,{0,220,255},2,cv::LINE_AA);
        for(int s=0;s<2;++s) {
            auto const& state=person.hands[s];auto const& k=person.selected;
            cv::Scalar arm=s==0?cv::Scalar(40,40,255):cv::Scalar(30,230,30);
            for(auto e:std::vector<std::pair<int,int>>{{5+s,7+s},{7+s,9+s}}) {
                segment(canvas,point(k[e.first]),point(k[e.second]),state.limb_consensus?arm:cv::Scalar(0,140,255),3,!state.limb_consensus);
                cv::circle(canvas,point(k[e.second]),4,arm,-1,cv::LINE_AA);
            }
            if(state.state=="rejected")continue;
            int start=s==0?91:112;
            auto hand=state.distance_eligible?(s==0?cv::Scalar(255,80,220):cv::Scalar(255,190,40)):cv::Scalar(0,140,255);
            for(int f=0;f<5;++f)for(int j=0;j<4;++j){int a=j==0?start:start+f*4+j;int b=start+1+f*4+j;segment(canvas,point(k[a]),point(k[b]),hand,2,!state.distance_eligible);}
            for(int j=0;j<21;++j)cv::circle(canvas,point(k[start+j]),3,hand,-1,cv::LINE_AA);
        }
    }
    for(auto const& h:r.hand_regions) {
        std::vector<cv::Point> points,hull;for(auto const& p:h.keypoints)points.push_back(point(p));cv::convexHull(points,hull);
        auto risk=std::find_if(r.risk.regions.begin(),r.risk.regions.end(),[&](auto const& item){return item.track_id==h.track_id&&item.side==h.side;});
        auto c=color(risk==r.risk.regions.end()?"unavailable":risk->risk_level);
        cv::polylines(canvas,std::vector<std::vector<cv::Point>>{hull},true,c,2,cv::LINE_AA);
    }
    for(auto const& region:r.risk.regions)if(region.source=="YOLO_BODY_BOX_FALLBACK") {
        auto it=std::find(r.track_ids.begin(),r.track_ids.end(),region.track_id);
        if(it!=r.track_ids.end()){auto const& b=r.boxes.at(it-r.track_ids.begin());cv::putText(canvas,"BODY FALLBACK",{int(b[0])+4,int(b[1])+50},cv::FONT_HERSHEY_SIMPLEX,.65,{0,160,255},2,cv::LINE_AA);}
    }
    cv::rectangle(canvas,{0,0},{canvas.cols,148},{20,24,30},-1);
    std::ostringstream headline;headline<<"RISK ";if(r.risk.risk_score)headline<<std::fixed<<std::setprecision(3)<<*r.risk.risk_score;else headline<<"N/A";headline<<"  "<<r.risk.risk_level;
    cv::putText(canvas,headline.str(),{24,40},cv::FONT_HERSHEY_SIMPLEX,1.0,color(r.risk.risk_level),2,cv::LINE_AA);
    std::ostringstream status;status<<"camera frame "<<r.frame_index<<" | seq "<<r.capture.driver_sequence<<" | latency "<<std::fixed<<std::setprecision(1)<<r.timing.capture_to_result_ms<<" ms";
    cv::putText(canvas,status.str(),{24,82},cv::FONT_HERSHEY_SIMPLEX,.72,{235,235,235},2,cv::LINE_AA);
    std::string calibration=r.door_mask_calibration_verified?"door calibration verified":"door calibration NOT verified (bring-up)";
    cv::putText(canvas,calibration+" | q/Esc: close preview",{24,123},cv::FONT_HERSHEY_SIMPLEX,.70,{0,190,255},2,cv::LINE_AA);
    return canvas;
}
void runLocalVisualization(CameraPipeline const& pipeline) {
    struct Window {
        Window(){cv::namedWindow(windowName,cv::WINDOW_NORMAL);cv::resizeWindow(windowName,1000,800);}
        ~Window(){try{if(cvGetWindowHandle(windowName))cv::destroyWindow(windowName);}catch(...) {}}
    } window;
    FrameResultPtr shown;
    while(pipeline.running()) {
        auto result=pipeline.latest();
        if(result&&result!=shown){cv::imshow(windowName,renderResult(*result));shown=std::move(result);}
        int key=cv::waitKey(16)&0xff;
        if(key==27||key=='q'||key=='Q'||cvGetWindowHandle(windowName)==nullptr)return;
    }
}
}