#include <cstdio>
#include "socket.h"
#include <mutex>
#include <string.h>
#include <sys/param.h>
#include <rpc/types.h>
#include <sys/socket.h>
#include <time.h>
#include <iostream>
#include <getopt.h>
#include <mutex>
#include <thread>
#include <unistd.h>
#include <time.h>
#include <vector>
#include <future>

#define MAX_REQUEST_TIME 512
#define MESSAGE_LEN 1024
#define PORTLEN 10
#define version 2.0
#define METHOD_GET 1
#define METHOD_HEAD 2
#define METHOD_TRACE 3
#define METHOD_OPTIONS 4
#define METHOD_POST 5
#define HTTP09 1
#define HTTP10 2
#define HTTP11 3
using namespace std;

int force=0;
int force_reload=0;
int clients=1;
int HTTP_PROTOCOL=HTTP11;
int METHOD=METHOD_GET;
int benchtime=100;
string PORT("80");
string HOST;
string REQUEST;
string PROXY;
mutex mut;
int num_c=0;
int bytes = 0;
int pages = 0;
int success = 0;
int failed = 0;
bool stop=false;

typedef struct ret_of_th
{
	int success_ = 0;
	int fail_ = 0;
	int bytes_ = 0;
} ret_of_th;


static void alarm_handler(int signal)
{
   stop = true;
}	

ret_of_th thr_socket(){
			char buffer[MESSAGE_LEN];
			int fail_ = 0;
			int success_ = 0;
			int bytes_ = 0;

			ret_of_th ret;

			while(stop == false){
				int fd = new_socket(PROXY.size() ? PROXY.data() : HOST.data(), PORT.data());
				if(fd == -1){ ++fail_; continue;; }

				memset(buffer, 0, sizeof(buffer));
				sprintf(buffer, "%s", REQUEST.data());
				int len = strlen(buffer);
				int wlen = write(fd, buffer, strlen(buffer));
				memset(buffer, 0, sizeof(buffer));

				if(force){
						if(len == wlen) ++success_;
						else ++fail_;
						close(fd);
						continue;
					}

				int rlen = read(fd, buffer, sizeof(buffer));

				// count speed need record time
				if(rlen > 0){ ++success_; bytes_ += rlen; }
				else if(rlen <= 0) ++fail_;
				close(fd);
			}

			ret.success_ = success_;
			ret.fail_ = fail_;
			ret.bytes_ = bytes_;

			return ret;
}

static const struct option long_options[]=
{
				{"force",no_argument,&force,1},
				{"reload",no_argument,&force_reload,1},
				{"http09",no_argument,&HTTP_PROTOCOL,HTTP09},
				{"http10",no_argument,&HTTP_PROTOCOL,HTTP10},
				{"http11",no_argument,&HTTP_PROTOCOL,HTTP11},
				{"get",no_argument,&METHOD,METHOD_GET},
				{"head",no_argument,&METHOD,METHOD_HEAD},
				{"trace",no_argument,&METHOD,METHOD_TRACE},
				{"options",no_argument,&METHOD,METHOD_OPTIONS},
				{"time",required_argument,NULL,'t'},
				{"clients",required_argument,NULL,'c'},
				{"proxy",required_argument,NULL,'p'},
				{"version",no_argument,NULL,'V'},
				{"help",no_argument,NULL,'?'},
				{0,0,0,0}
};

static void usage(void)
{
				std::cout<<" webbench [option]...URL\n"
								<<"  -f|--force                Don't wait for reply from server.\n"
								<<"  -r|--reload               Send reload request - Pragma: no-cache.\n"
								<<"  -t|--time   <sec>         Run benchmark for <sec> seconds. Default 30.\n"
								<<"  -p|--proxyy <server:port> Use proxy server for requeset.\n"
								<<"  -c|--clients <n> Run <n>  HTTP clients at once.. Deafult one.\n"
								<<"  -9|--http09               Use HTTP/0.9 style requests.\n"
								<<"  -1|--http10               Use HTTP/1.0 protocol.\n"
								<<"  -2|--http11               Use HTTP/1.1 protocol.\n"
								<<"  --get                     Use GET requeset method.\n"
								<<"  --head                    Use HEAD request method.\n"
								<<"  --options                 Use OPTIONS request method.\n"
								<<"  --trace                   Use TRACE request method.\n"
								<<"  -?|-h|--help              This information.\n"
								<<"  -V|--version              Display program version.\n";
}

static void show_ver(void){
				std::cout<<"version:"<<version<<"\nwrite by Wlzzzz-del"<<std::endl;
}

static void welcome(void){
				cout<<"WebBench - Simple Web Benchmark 1.5"<<endl;
				cout<<"Copyright(c) Radim Kolar 1997-2004, GPL Open Source Software."<<endl;
				cout<<"Overwrite by Lyu Chuan 2023.---http://www.github.com/Rubint0x/"<<endl;
				cout<<endl<<"Request:"<<endl;
}

int buildrequest(char* _url){
				string url(_url);
				string temp;
				const int size = url.size();
				if(size > 1500){cout<<"url too long"<<endl;return -1;}
				if(size < 8){cout<<"url is too short"<<endl;return -1;}
				if(url.find("http://")){cout<<"not found http://"<<endl;return -1;}

				temp=url.substr(url.find("http://")+7, size - 8);
				//find返回第一次查找到该子串的下表位置即0，+7为将http://切割掉，substr的第二个参数是往后提取多少位，即size-7再将“/”减掉就是-8。

				/// find port
				if(temp.find(":") != string::npos)
				PORT=(temp.substr(temp.find(":")+1,temp.size()-temp.find(":")));

				// find host
				HOST= temp.substr(0,temp.find(":"));

				// assemble request
				switch (METHOD) {
								case METHOD_GET:{REQUEST += "GET"; break; }
								case METHOD_HEAD:{REQUEST += "HEAD"; break; }
								case METHOD_OPTIONS:{REQUEST += "OPTIONS"; break; }
								case METHOD_TRACE:{REQUEST += "TRACE"; break; }
				}

				REQUEST += " http://";
				REQUEST += HOST;
				switch (HTTP_PROTOCOL) {
								case HTTP09:{REQUEST+= " HTTP/0.9"; break;}
								case HTTP10:{REQUEST+= " HTTP/1.0"; break;}
								case HTTP11:{REQUEST+= " HTTP/1.1"; break;}
				}
				REQUEST+= "\r\nUser-Agent: WebBench 2.0\r\n";
				REQUEST+= "Host: ";
				REQUEST+= HOST;
				if(force_reload && PROXY.size())
				REQUEST+= "\r\nParagma: no-cache";
				REQUEST+="\r\nConnection: close\r\n\r\n";
				return 0;
}

int main(int argc, char *argv[]){
		int opt=0;
		int options_index=0;
		char *tmp=NULL;

		if(argc == 1){
						usage();
						return 2;
		}
		while(((opt=getopt_long(argc,argv,"frt:p:c:912?V",long_options,&options_index))!=EOF)){
						switch(opt){
										case 0: break;
										case 'f':{force=1;break;}
										case 'r':{force_reload=1;break;}
										case 't':{benchtime=atoi(optarg);break;};
										case 'p':{
													// get port and PROXY
													string temp(optarg);
													int loc=temp.find(':');
													if(loc==std::string::npos){std::cout<<"Not found proxy ,you need to specify after :";return -1;}
													PROXY = temp.substr(0,loc);
													PORT = temp.substr(loc+1,temp.size()-loc);
													std::cout<<PORT;
													break;
												}
										case 'c':{clients=atoi(optarg);break;};
										case '9':{HTTP_PROTOCOL=HTTP09;break;};
										case '1':{HTTP_PROTOCOL=HTTP10;break;};
										case '2':{HTTP_PROTOCOL=HTTP11;break;};
										case '?':{usage();return 0;};
										case 'V':{show_ver();return 0;};
						}
		}

						if(optind==argc){
										cout<<"not found url"<<endl;
										return -1;
						}
						if(HTTP_PROTOCOL==HTTP09) METHOD=METHOD_GET;

						if(METHOD==METHOD_TRACE || METHOD==METHOD_OPTIONS) HTTP_PROTOCOL=HTTP11;

						if(buildrequest(argv[optind])==-1)return -1;
						welcome();
						cout<<REQUEST;

 						struct sigaction sa;
						sa.sa_handler=alarm_handler;
						sa.sa_flags=0;
						if(sigaction(SIGALRM, &sa, NULL)) exit(3);
						alarm(benchtime);

						std::cout<<"Running info:";
						std::cout<<clients<<" clients, running "<<benchtime<<" sec "<<endl;
						std::vector< std::future<ret_of_th> > th;
						for(num_c=0; num_c < clients; ++num_c){
								th.push_back(std::async(thr_socket));	
						}

						for(num_c=0; num_c < clients; ++num_c){
								ret_of_th ret = th[num_c].get();
								success += ret.success_;
								failed += ret.fail_;
								bytes += ret.bytes_;
						}


						std::cout<<"Speed= "<<(float)(success+failed)/((float)benchtime/60.0)<<" pages/min "<<bytes/benchtime<<" bytes/second "<<endl;
						std::cout<<"Request: "<<success<<" susceed, "<<failed<<" failed"<<endl;
						return 0;
}