#include "llmchatcompletionrequest.hpp"

#include <nlohmann/json.hpp>

namespace ecode {

LLMChatCompletionRequest::LLMChatCompletionRequest( const std::string& uri, const std::string& auth,
													const std::string& reqBody,
													const std::string& provider ) :
	mUrl( uri ), mHttp( Http::Pool::getGlobal().get( URI( uri ) ) ) {
	mRequest.setUri( mUrl.getPathEtc() );
	mRequest.setHeader( "Content-Type", "application/json" );
	if ( provider == "anthropic" ) {
		mRequest.setHeader( "x-api-key", auth );
		mRequest.setHeader( "anthropic-version", "2023-06-01" );
	} else if ( !auth.empty() ) {
		mRequest.setHeader( "Authorization", String::format( "Bearer %s", auth ) );
	}
	mRequest.setBody( reqBody );
	mRequest.setFollowRedirect( true );
	mRequest.setMethod( Http::Request::Method::Post );
	mRequest.setCancelCallback( [this]( const Http&, const Http::Request& ) { onCancel(); } );
	mRequest.setProgressCallback( [this]( const Http&, const Http::Request&, const Http::Response&,
										  const Http::Request::Status& status, size_t, size_t ) {
		if ( mCancel ) {
			onCancel();
			return false;
		}
		mHadProgress = true;
		if ( status != Http::Request::ContentReceived )
			return true;
		std::string chunk =
			mReadBytes ? mStream.getStream().substr( mReadBytes ) : mStream.getStream();
		mReadBytes = mStream.getStream().size();

		mBuffer += chunk;
		auto lastNL = mBuffer.find_last_of( '\n' );

		if ( lastNL == std::string::npos )
			return true;

		std::string_view processChunk = std::string_view{ mBuffer }.substr( 0, lastNL );

		String::readBySeparator( processChunk, [this]( std::string_view subchunk ) {
			if ( subchunk.empty() )
				return;
			// OpenAI format?
			if ( !String::startsWith( subchunk, "data: " ) ) {
				// ollama format?
				if ( !subchunk.empty() && subchunk[0] == '{' ) {
					nlohmann::json j = nlohmann::json::parse( subchunk.begin(), subchunk.end(),
															  nullptr, false, true );

					if ( j.contains( "message" ) && j["message"].contains( "content" ) ) {
						const auto& msg = j["message"];
						std::string delta = msg["content"];
						if ( delta.empty() )
							return;
						if ( mFirstMessage ) {
							if ( String::startsWith( delta, "<think>" ) )
								mReasoning = true;
							mFirstMessage = false;
						}

						if ( streamedResponseCb )
							streamedResponseCb( delta, mReasoning );
						if ( mReasoning ) {
							if ( String::endsWith( delta, "</think>" ) )
								mReasoning = false;
							mReasoningResponse += std::move( delta );
						} else {
							mResponse += std::move( delta );
						}
					}
				}

				return;
			}
			std::string jsonStr( subchunk.substr( 6 ) );
			if ( jsonStr == "[DONE]" )
				return;
			nlohmann::json data = nlohmann::json::parse( jsonStr, nullptr, false, true );

			// OpenAI
			if ( data.contains( "choices" ) && data["choices"].is_array() ) {
				nlohmann::json& choices = data["choices"];

				for ( const auto& choice : choices ) {
					if ( choice["delta"].contains( "reasoning_content" ) ) {
						const auto& reasoningContent = choice["delta"]["reasoning_content"];
						if ( reasoningContent.is_string() ) {
							std::string delta = choice["delta"]["reasoning_content"];
							if ( !delta.empty() ) {
								if ( streamedResponseCb )
									streamedResponseCb( delta, true );
								mReasoningResponse += std::move( delta );
							}
						}
					} else if ( choice["delta"].contains( "reasoning" ) ) {
						const auto& reasoningContent = choice["delta"]["reasoning"];
						if ( reasoningContent.is_string() ) {
							std::string delta = choice["delta"]["reasoning"];
							if ( !delta.empty() ) {
								if ( streamedResponseCb )
									streamedResponseCb( delta, true );
								mReasoningResponse += std::move( delta );
							}
						}
					} else if ( choice["delta"].contains( "content" ) ) {
						const auto& content = choice["delta"]["content"];
						if ( content.is_string() ) {
							std::string delta = choice["delta"]["content"];
							if ( !delta.empty() ) {
								if ( streamedResponseCb )
									streamedResponseCb( delta, false );
								mResponse += std::move( delta );
							}
						}
					}
				}
				// Anthropic
			} else if ( data.contains( "delta" ) &&
						data["delta"].value( "type", "" ) == "text_delta" ) {
				std::string delta = data["delta"]["text"];
				if ( !delta.empty() ) {
					if ( streamedResponseCb )
						streamedResponseCb( delta, false );
					mResponse += std::move( delta );
				}
			}
		} );

		mBuffer.erase( 0, lastNL + 1 );
		return true;
	} );
}

LLMChatCompletionRequest::~LLMChatCompletionRequest() {
	cancelCb = nullptr;
	doneCb = nullptr;
	streamedResponseCb = nullptr;
	cancel( true );
}

void LLMChatCompletionRequest::request() {
	mCancel = false;
	mCancelled = false;
	mHadProgress = false;
	mBuffer.clear();
	Http::Response res = mHttp->downloadRequest( mRequest, mStream, Seconds( 5 ) );
	if ( doneCb )
		doneCb( *this, res );
}

void LLMChatCompletionRequest::requestAsync() {
	mCancel = false;
	mCancelled = false;
	mHadProgress = false;
	mBuffer.clear();
	mRequestId = mHttp->downloadAsyncRequest(
		[this]( const Http&, Http::Request&, Http::Response& res ) {
			if ( doneCb )
				doneCb( *this, res );
		},
		mRequest, mStream, Seconds( 5 ) );
}

void LLMChatCompletionRequest::cancel(bool resetCancelCallback ) {
	mCancel = true;
	if ( mRequestId )
		mHttp->setCancelRequest( mRequestId, resetCancelCallback );
	if ( !mHadProgress )
		onCancel();
}

const std::string& LLMChatCompletionRequest::getStream() const {
	return mStream.getStream();
}

bool LLMChatCompletionRequest::isCancelled() const {
	return mCancelled;
}

void LLMChatCompletionRequest::onCancel() {
	if ( !mCancelled && cancelCb )
		cancelCb( *this );
	mCancelled = true;
}

} // namespace ecode
