#include "effects/ConeWave.h"

#include <osg/BlendFunc>
#include <osg/Material>
#include <osg/Texture2D>
#include <osg/ShapeDrawable>
#include <osg/Group>
#include <osgViewer/Viewer>
#include <osgDB/ReadFile>
#include <osg/Shader>
#include <osg/Program>
#include <osg/Uniform>
#include <osg/Depth>
#include <osg/Cullface>
#include <osgEarth/Registry>

class WaveSurfaceCallback : public osg::NodeCallback {
public:
    WaveSurfaceCallback() : timeElapsed(0.0f) {}

    virtual void operator()(osg::Node* node, osg::NodeVisitor* nv) {
        timeElapsed += (nv->getFrameStamp()->getSimulationTime() - SimulationTime);

        osg::Geode* geode = dynamic_cast<osg::Geode*>(node);
        if (geode) {
            osg::Geometry* geometry = dynamic_cast<osg::Geometry*>(geode->getDrawable(0));
            if (geometry) {
                osg::Vec3Array* vertices = dynamic_cast<osg::Vec3Array*>(geometry->getVertexArray());
                if (vertices) {
                    float amplitude = 2.0f;
                    float frequency = 1.0f;
                    for (unsigned int i = 0; i < vertices->size(); ++i) {
                        osg::Vec3f& vertex = (*vertices)[i];
                        vertex.z() = amplitude * sin(frequency * (vertex.x() + timeElapsed));
                    }
                    geometry->setVertexArray(vertices);
                    geometry->dirtyBound();
                }
            }
        }

        traverse(node, nv);
        SimulationTime = nv->getFrameStamp()->getSimulationTime();
    }

private:
    float timeElapsed;
    double SimulationTime;
};

class RadarWaveTimeCallback : public osg::NodeCallback {
public:
    RadarWaveTimeCallback(osg::ref_ptr<osg::Uniform> waveTimeUniform) 
        : waveTimeUniform_(waveTimeUniform), startTime_(0.0) {}

    virtual void operator()(osg::Node* node, osg::NodeVisitor* nv) {
        if (nv->getFrameStamp()) {
            double currentTime = nv->getFrameStamp()->getSimulationTime();
            if (startTime_ == 0.0) {
                startTime_ = currentTime;
            }
            
            float elapsedTime = static_cast<float>(currentTime - startTime_);
            
            if (waveTimeUniform_.valid()) {
                waveTimeUniform_->set(elapsedTime);
            }
        }
        
        traverse(node, nv);
    }

private:
    osg::ref_ptr<osg::Uniform> waveTimeUniform_;
    double startTime_;
};


ConeWave::ConeWave() {
    osgEarth::Registry::shaderGenerator().run(this);
    currentTime_ = 0.0;
    
    height_ = 1000.0f;
    radius_ = 50.0f;
    waveRadius_ = 100.0f;
    waveSpeed_ = 20.0f;
    waveCount_ = 3;
    baseColor_ = osg::Vec4(0.0f, 0.8f, 1.0f, 1.0f);
    waveColor_ = osg::Vec4(1.0f, 0.5f, 1.0f, 0.8f);
    
    ringBrightAlpha_ = 0.8f;
    ringDarkAlpha_ = 0.3f;
    coneAlpha_ = 0.7f;
}


ConeWave::~ConeWave(void)
{
}

void ConeWave::Render(double dt) {
}

void ConeWave::InitGeode() {
    Rebuild();
    getOrCreateStateSet()->setMode(GL_LIGHTING, osg::StateAttribute::OFF);
    getOrCreateStateSet()->setMode(GL_BLEND, osg::StateAttribute::ON);
    getOrCreateStateSet()->setRenderingHint(osg::StateSet::TRANSPARENT_BIN);
    getOrCreateStateSet()->setMode(GL_DEPTH_TEST, osg::StateAttribute::ON);
    setCullingActive(true);
}

void ConeWave::Destory() {
    OnDestroy();
}

void ConeWave::SetHeight(float height) {
    height_ = height;
    
	Rebuild();
}

void ConeWave::SetBaseColor(const osg::Vec4& color) {
    baseColor_ = color;

    if (baseColorUniform_.valid()) {
        baseColorUniform_->set(baseColor_);
    }
}

void ConeWave::SetWaveCount(int count) {
    waveCount_ = count;
    
    if (waveCountUniform_.valid()) {
        waveCountUniform_->set(static_cast<float>(waveCount_));
    } else {
        Rebuild();
    }
}

void ConeWave::SetWaveSpeed(float speed) {
    waveSpeed_ = speed;
 
    if (waveSpeedUniform_.valid()) {
        waveSpeedUniform_->set(waveSpeed_);
    }
}

void ConeWave::SetWaveRadius(float radius) {
    waveRadius_ = radius;
    radius_ = radius;
    Rebuild();
}

void ConeWave::SetWaveColor(const osg::Vec4& color) {
    waveColor_ = color;

    if (waveColorUniform_.valid()) {
        waveColorUniform_->set(waveColor_);
    }
}

void ConeWave::SetRingBrightAlpha(float alpha) {
    ringBrightAlpha_ = alpha;

    if (ringBrightAlphaUniform_.valid()) {
        ringBrightAlphaUniform_->set(alpha);
    }
}

void ConeWave::SetRingDarkAlpha(float alpha) {
    ringDarkAlpha_ = alpha;
    if (ringDarkAlphaUniform_.valid()) {
        ringDarkAlphaUniform_->set(alpha);
    }
}

void ConeWave::SetConeAlpha(float alpha) {
    coneAlpha_ = alpha;
    if (coneGeometry_.valid()) {
        osg::StateSet* ss = coneGeometry_->getOrCreateStateSet();
        osg::ref_ptr<osg::Material> material = new osg::Material();
        material->setDiffuse(osg::Material::FRONT_AND_BACK, osg::Vec4(0.0f, 0.8f, 1.0f, alpha));
        material->setTransparency(osg::Material::FRONT_AND_BACK, 1.0f - alpha);
        ss->setAttributeAndModes(material, osg::StateAttribute::ON);
    }
}

void ConeWave::Clear() {
    if (coneCallback_.valid()) {
        setUpdateCallback(nullptr);
        coneCallback_ = nullptr;
	}

    if (coneGeometry_.valid()) {
        removeDrawable(coneGeometry_);
        coneGeometry_ = nullptr;
        coneDrawable_ = nullptr;
    }
    
    waveTimeUniform_ = nullptr;
    baseColorUniform_ = nullptr;
    waveColorUniform_ = nullptr;
    levelHeightUniform_ = nullptr;
    waveSpeedUniform_ = nullptr;
    ringBrightAlphaUniform_ = nullptr;
	ringDarkAlphaUniform_ = nullptr;
}

void ConeWave::Rebuild() {
    Clear();
    
    CreateRadarScanWave();
}

void ConeWave::CreateConeGeometry() {
    osg::ref_ptr<osg::Geometry> coneGeometry = new osg::Geometry();
    
    const int segments = 64;
    const float angleStep = 2.0f * osg::PI / segments;
    
    osg::ref_ptr<osg::Vec3Array> vertices = new osg::Vec3Array();
    osg::ref_ptr<osg::Vec3Array> normals = new osg::Vec3Array();
    osg::ref_ptr<osg::Vec2Array> texCoords = new osg::Vec2Array();
    
    vertices->push_back(osg::Vec3(0.0f, 0.0f, 0.0f));
    normals->push_back(osg::Vec3(0.0f, 0.0f, -1.0f));
    texCoords->push_back(osg::Vec2(0.5f, 0.5f));
    
    for (int i = 0; i <= segments; ++i) {
        float angle = i * angleStep;
        float x = radius_ * cos(angle);
        float y = radius_ * sin(angle);
        float z = height_;
        
        vertices->push_back(osg::Vec3(x, y, z));
        
        osg::Vec3 normal(x, y, radius_);
        normal.normalize();
        normals->push_back(normal);
        
        texCoords->push_back(osg::Vec2((x/radius_ + 1.0f) * 0.5f, (y/radius_ + 1.0f) * 0.5f));
    }
    
    coneGeometry->setVertexArray(vertices);
    coneGeometry->setNormalArray(normals);
    coneGeometry->setNormalBinding(osg::Geometry::BIND_PER_VERTEX);
    coneGeometry->setTexCoordArray(0, texCoords);
    
    osg::ref_ptr<osg::DrawElementsUInt> indices = new osg::DrawElementsUInt(osg::PrimitiveSet::TRIANGLES);
    
    for (int i = 0; i < segments; ++i) {
        indices->push_back(0);
        indices->push_back(i + 1);
        indices->push_back(i + 2);
    }
    
    coneGeometry->addPrimitiveSet(indices);
    if (coneGeometry_.valid() && coneGeometry_->getNumParents() > 0) {
        removeDrawable(coneGeometry_);
        coneGeometry_ = coneGeometry;
        coneDrawable_ = coneGeometry_.get();
        waveTimeUniform_ = nullptr;
        levelHeightUniform_ = nullptr;
        addDrawable(coneDrawable_);
    } else {
        coneGeometry_ = coneGeometry;
        coneDrawable_ = coneGeometry_.get();
    }
}

void ConeWave::CreateRadarScanWave() {
    CreateConeGeometry();
    
    addDrawable(coneDrawable_);
    
    osg::StateSet* ss = coneDrawable_->getOrCreateStateSet();
    ss->setRenderBinDetails(120, "RenderBin");
    ss->setMode(GL_CULL_FACE, osg::StateAttribute::OFF);
    ss->setMode(GL_LIGHTING, osg::StateAttribute::OFF);
    
    osg::ref_ptr<osg::BlendFunc> bf = new osg::BlendFunc(osg::BlendFunc::SRC_ALPHA, osg::BlendFunc::ONE_MINUS_SRC_ALPHA);
    ss->setAttributeAndModes(bf, osg::StateAttribute::ON);
    ss->setMode(GL_BLEND, osg::StateAttribute::ON);
    ss->setRenderingHint(osg::StateSet::TRANSPARENT_BIN);
    ss->setAttributeAndModes(new osg::CullFace(osg::CullFace::BACK));
    
    CreateRadarShader();
    
    osg::ref_ptr<osg::NodeCallback> waveTimeCallback = new RadarWaveTimeCallback(waveTimeUniform_);
    coneCallback_ = waveTimeCallback;
    setUpdateCallback(coneCallback_);
}

void ConeWave::CreateRadarShader() {
    if (waveTimeUniform_.valid()) {
        return;
    }
    
    static const char* vertexShaderSource = 
        "varying vec3 pos;\n"
        "void main()\n"
        "{\n"
        "    pos.x = gl_Vertex.x;\n"
        "    pos.y = gl_Vertex.y;\n"
        "    pos.z = gl_Vertex.z;\n"
        "    gl_Position = gl_ModelViewProjectionMatrix * gl_Vertex;\n"
        "}\n";
    
    static const char* fragmentShaderSource = 
        "uniform float height;\n"
        "uniform vec4 baseColor;\n"
        "uniform vec4 waveColor;\n"
        "uniform float waveTime;\n"
        "uniform float ringBrightAlpha;\n"
        "uniform float ringDarkAlpha;\n"
        "uniform float waveSpeed;\n"
        "uniform float waveCount;\n"
        "varying vec3 pos;\n"
        "void main()\n"
        "{\n"
        "    float h = abs(pos.z) / max(height, 1.0);\n"
        "    float radialDist = sqrt(pos.x * pos.x + pos.y * pos.y);\n"
        "    float waveFreq = 0.2 * max(1.0, waveCount * 0.1);\n"
        "    float timeScale = waveTime * waveSpeed * 0.1;\n"
        "    float wavePhase = radialDist * waveFreq - timeScale;\n"
        "    float ripple = sin(wavePhase);\n"
        "    float ripple2 = sin(wavePhase * 1.1 + timeScale * 0.5);\n"
        "    ripple = (ripple + ripple2 * 0.3) / 1.3;\n"
        "    if (ripple > 0.3)\n"
        "    {\n"
        "        gl_FragColor = vec4(waveColor.rgb, ringBrightAlpha);\n"
        "    }\n"
        "    else\n"
        "    {\n"
        "        gl_FragColor = vec4(baseColor.rgb, ringDarkAlpha);\n"
        "    }\n"
        "}\n";
   
    osg::ref_ptr<osg::Shader> vertexShader = new osg::Shader(osg::Shader::VERTEX, vertexShaderSource);
    osg::ref_ptr<osg::Shader> fragmentShader = new osg::Shader(osg::Shader::FRAGMENT, fragmentShaderSource);
    
    osg::ref_ptr<osg::Program> program = new osg::Program;
    program->addShader(vertexShader);
    program->addShader(fragmentShader);
    
    osg::StateSet* stateSet = coneDrawable_->getOrCreateStateSet();
    stateSet->setAttributeAndModes(program, osg::StateAttribute::ON);
    
    waveTimeUniform_ = new osg::Uniform("waveTime", 0.0f);
    baseColorUniform_ = new osg::Uniform("baseColor", baseColor_);
    waveColorUniform_ = new osg::Uniform("waveColor", waveColor_);
    levelHeightUniform_ = new osg::Uniform("height", height_ > 0 ? height_ : 100.0f);
    waveSpeedUniform_ = new osg::Uniform("waveSpeed", waveSpeed_ > 0 ? waveSpeed_ : 1.0f);
    waveCountUniform_ = new osg::Uniform("waveCount", waveCount_ > 0 ? waveCount_ : 1);
    
    ringBrightAlphaUniform_ = new osg::Uniform("ringBrightAlpha", ringBrightAlpha_);
    ringDarkAlphaUniform_ = new osg::Uniform("ringDarkAlpha", ringDarkAlpha_);
    
    stateSet->addUniform(waveTimeUniform_);
    stateSet->addUniform(baseColorUniform_);
    stateSet->addUniform(waveColorUniform_);
    stateSet->addUniform(levelHeightUniform_);
    stateSet->addUniform(waveSpeedUniform_);
    stateSet->addUniform(waveCountUniform_);
    stateSet->addUniform(ringBrightAlphaUniform_);
    stateSet->addUniform(ringDarkAlphaUniform_);
}