package com.microsoft.graph.httpcore;

import java.net.URI;
import java.net.URISyntaxException;

import org.apache.http.Header;
import org.apache.http.HttpRequest;
import org.apache.http.HttpResponse;
import org.apache.http.HttpStatus;
import org.apache.http.ProtocolException;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpHead;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.client.methods.RequestBuilder;
import org.apache.http.impl.client.DefaultRedirectStrategy;
import org.apache.http.protocol.HttpContext;
import org.apache.http.util.Args;

public class RedirectHandler extends DefaultRedirectStrategy{
	
	public static final RedirectHandler INSTANCE = new RedirectHandler();
	
    @Override
    public boolean isRedirected(
            final HttpRequest request,
            final HttpResponse response,
            final HttpContext context) throws ProtocolException {
        Args.notNull(request, "HTTP request");
        Args.notNull(response, "HTTP response");

        final int statusCode = response.getStatusLine().getStatusCode();
        final Header locationHeader = response.getFirstHeader("location");
        if(locationHeader == null)
        	return false;
        
        if(statusCode == HttpStatus.SC_MOVED_TEMPORARILY ||
        		statusCode == HttpStatus.SC_MOVED_PERMANENTLY ||
        		statusCode == HttpStatus.SC_TEMPORARY_REDIRECT ||
        		statusCode == HttpStatus.SC_SEE_OTHER)
        	return true;
        
        return false;
    }
    
    @Override
    public HttpUriRequest getRedirect(
            final HttpRequest request,
            final HttpResponse response,
            final HttpContext context) throws ProtocolException {
        final URI uri = getLocationURI(request, response, context);
        final String method = request.getRequestLine().getMethod();
        if (method.equalsIgnoreCase(HttpHead.METHOD_NAME)) {
            return new HttpHead(uri);
        } else if (method.equalsIgnoreCase(HttpGet.METHOD_NAME))
            return new HttpGet(uri);
        else {
            final int status = response.getStatusLine().getStatusCode();
            if(status != HttpStatus.SC_SEE_OTHER) {
            	try {
            		final URI requestURI = new URI(request.getRequestLine().getUri());
            		if(!uri.getHost().equalsIgnoreCase(requestURI.getHost()) || 
            				!uri.getScheme().equalsIgnoreCase(requestURI.getScheme()))
            			request.removeHeaders("Authorization");
            		return RequestBuilder.copy(request).setUri(uri).build();
            	}
            	 catch (final URISyntaxException ex) {
                     throw new ProtocolException(ex.getMessage(), ex);
                 }
            }
            return new HttpGet(uri);
        }
    }
}

@@ -1601,32 +1601,26 @@
 direct(%0A    
-        
+%09%09
 final HttpRe
@@ -1634,32 +1634,26 @@
 equest,%0A    
-  
-      
+%09%09
 final HttpRe
@@ -1669,32 +1669,26 @@
 sponse,%0A    
-        
+%09%09
 final HttpCo
@@ -1729,28 +1729,25 @@
 ption %7B%0A    
-    
+%09
 final URI ur
@@ -1798,425 +1798,8 @@
 t);%0A
-        final String method = request.getRequestLine().getMethod();%0A        if (method.equalsIgnoreCase(HttpHead.METHOD_NAME)) %7B%0A            return new HttpHead(uri);%0A        %7D else if (method.equalsIgnoreCase(HttpGet.METHOD_NAME))%0A            return new HttpGet(uri);%0A        else %7B%0A            final int status = response.getStatusLine().getStatusCode();%0A            if(status != HttpStatus.SC_SEE_OTHER) %7B%0A        
     
@@ -1805,24 +1805,16 @@
  %09try %7B%0A
-        
     %09%09fi
@@ -1878,24 +1878,16 @@
 Uri());%0A
-        
     %09%09if
@@ -1945,24 +1945,16 @@
 ()) %7C%7C %0A
-        
     %09%09%09%09
@@ -2012,24 +2012,16 @@
 eme()))%0A
-        
     %09%09%09r
@@ -2062,109 +2062,22 @@
 n%22);
-%0A            %09%09return RequestBuilder.copy(request).setUri(uri).build();%0A            %09%7D%0A            %09 
+%09%0A    %09%7D%0A    %09
 catc
@@ -2114,33 +2114,18 @@
 ) %7B%0A    
-   
-              
+%09%09
 throw ne
@@ -2174,49 +2174,132 @@
     
-             
+%09
 %7D%0A    
+%09%0A
     
-    %7D%0A        
+%09final int status = response.getStatusLine().getStatusCode();%0A    %09if(status == HttpStatus.SC_SEE_OTHER)%0A
     
+%09%09
 retu
@@ -2323,21 +2323,73 @@
 i);%0A    
-    %7D
+%09return RequestBuilder.copy(request).setUri(uri).build();
 %0A    %7D%0A%7D

Update getRedirect logic in RedirectHandler