#include <netinet/in.h>
#include "tcp.h"

using namespace std;

Tcp::Tcp(int sock) {
	sockfd = sock;
	dataType = ASCII;
	localAddr = 0;
	localPort = 0;
	eof = false;
}

Tcp::~Tcp() {
	if(close(sockfd)) throw "Close socket failed";
}

void Tcp::setDataType(DataType type) {
	dataType = type;
}

string Tcp::read() {
	switch(dataType) {
	case ASCII:
		buffer.erase(0, pos+2);
		pos = 1;
		while(true) {
			pos = buffer.find("\r\n", pos-1);
			if(pos!=string::npos)
				return buffer.substr(0, pos);
			else
				pos = buffer.length();

			char buf[1024];
			int size = ::read(sockfd, buf, 1024);

			if(size>0) {
				buffer.append(buf, size);
			} else if(size==0) {
				eof = true;
				return buffer;
			} else
				throw "Read socket failed";
		}

	case IMAGE: {
		char buf[1024];
		int size = ::read(sockfd, buf, 1024);

		if(size==0)
			eof = true;
		else if(size<0)
			throw "Read socket failed";

		return string(buf, size);
	}

	default:
		throw "Unsupported data type";
	}
}

void Tcp::readToStream(ostream& os) {
	while(true) {
		os << read();
		if(eof) return;

		if(dataType==ASCII) os << endl;
	}
}

void Tcp::write(const string& buf) {
	if(::write(sockfd, buf.c_str(), buf.size())==-1) throw "Write socket failed";
	if(dataType==ASCII)
		if(::write(sockfd, "\r\n", 2)==-1) throw "Write socket failed";
}

void Tcp::writeFromStream(istream& is) {
	switch(dataType) {
	case ASCII: {
		string text;
		while(getline(is, text)) write(text);

		break;
	}

	case IMAGE: {
		char buf[1024];
		while(is) {
			is.read(buf, 1024);
			write(string(buf, is.gcount()));
		}

		break;
	}

	default:
		throw "Unsupported data type";
	}
}

int Tcp::getLocalAddr() {
	if(localAddr==0) {
		sockaddr_in addr;
		socklen_t len = sizeof addr;

		if(getsockname(sockfd, (sockaddr*)&addr, &len))
			throw "Get local address failed";

		localAddr = ntohl(addr.sin_addr.s_addr);
		localPort = ntohs(addr.sin_port);
	}

	return localAddr;
}

int Tcp::getLocalPort() {
	if(localPort==0) {
		sockaddr_in addr;
		socklen_t len = sizeof addr;

		if(getsockname(sockfd, (sockaddr*)&addr, &len))
			throw "Get local port failed";

		localAddr = ntohl(addr.sin_addr.s_addr);
		localPort = ntohs(addr.sin_port);
	}

	return localPort;
}