import groovyx.net.http.HTTPBuilder;

import java.io.IOException;
import java.net.URISyntaxException;
import java.util.HashMap;
import java.util.Map;

import org.apache.http.client.ClientProtocolException;
import org.scribe.http.Request;
import org.scribe.http.Request.Verb;
import org.scribe.oauth.Scribe;
import org.scribe.oauth.Token;

/**
 * Adds OAuth Consumer support to HttpBuilder using Pablo Fernandez's Scribe
 * library (http://github.com/fernandezpablo85/scribe).
 * 
 * Instantiate your own Scribe and pass it to the <code>create</code> factory
 * method. Provide an <code>accessToken</code> parameter in the
 * <code>args</code> map of each request or use <code>setAccessToken</code> to
 * use the same one repeatedly.
 * 
 * Example: <code>
 * def serviceProps = new Properties([    "consumer.key" : "myApiToken", 
 *                                     "consumer.secret" : "myApiSecret" ])
 * 
 * def scribe = new Scribe(serviceProps)
 * 
 * //Use Scribe as described on its github to get the request and access 
 * //tokens. We'll pretend we just have one.
 * 
 * def access = new Token("theAccessToken","andItsSecret")
 * 
 * def http = OAuthHttpBuilder.create(
 *               "http://www.twitter.com/oauth/some_resource", 
 *               scribe 
 *            );
 * 
 * http.get(      query : [ arg1 : "value1", arg2 : "value2" ], 
 *           accessToken: access );
 * </code>
 * 
 * @author Erem Boto
 */
@SuppressWarnings("unchecked")
public class OAuthHttpBuilder {
  private final Scribe signer;
  private final HTTPBuilder delegate;
  private Token defaultAccessToken = null;

  
  /*
   * Class provision
   */
  public static OAuthHttpBuilder create( Object defaultUri,
                                         Object defaultContentType, 
                                         Scribe signer ) throws URISyntaxException {
    return new OAuthHttpBuilder(
        new HTTPBuilder(defaultUri, defaultContentType), signer);
  }

  public static OAuthHttpBuilder create( Object defaultUri, 
                                         Scribe signer )
  throws URISyntaxException {
    return new OAuthHttpBuilder(new HTTPBuilder(defaultUri), signer);
  }

  public OAuthHttpBuilder(HTTPBuilder delegate, Scribe signer) {
    this.signer = signer;
    this.delegate = delegate;
  }

  /*
   * Public Methods
   */
  /**
   * Sets the default access token to use when making OAuth requests.
   * 
   * @param token
   */
  public void setAccessToken(Token token) {
    this.defaultAccessToken = token;
  }

  /*
   * Wrapped HttpBuilder Methods
   */
  public Object get(Map<String, ?> args) 
  throws ClientProtocolException, IOException, URISyntaxException {
    preProcessArgsHeader(Verb.GET, args);
    return delegate.get(args);
  }

  public Object post(Map<String, ?> args) 
  throws ClientProtocolException, URISyntaxException, IOException {
    preProcessArgsHeader(Verb.POST, args);
    return delegate.post(args);
  }

  
  /*
   * Private Methods
   */
  /**
   * Uses provided "accessToken" and "uri" parameters to construct an
   * Authorization header with Scribe.
   * 
   * @param method
   *          the request's HTTP method
   * @param args
   *          the argument map as passed to get() or post()
   */
  public void preProcessArgsHeader(Verb method, Map<String, ?> args) {
    Token accessToken = accessTokenFromArgs(args);
    String uriParam = uriParamFromArgs(args);

    // God knows what happens if they aren't strings.
    Map<String, String> query = (Map<String, String>) args.get("query");

    Object bodyParams = args.get("body");

    // because we're using one map, the request won't be to spec if you have
    // repeated query parameters =(
    Map oauthParams = new HashMap();
    if (query != null) {
      oauthParams.putAll(query);
    }

    boolean shouldPutBody = (method == Verb.POST && bodyParams instanceof Map);
    if (shouldPutBody) {
      oauthParams.putAll((Map<String, String>) bodyParams);
    }

    String authHeader = makeAuthHeader( method, 
                                        uriParam,
                                        accessToken,
                                        oauthParams );
    addAuthorizationHeader(authHeader, args);
  }

  /**
   * Grabs the Scribe access token from the provided arguments, or returns the
   * default token.
   * 
   * @param args
   *          the arguments passed to either get or post.
   * @return the accessToken that was provided
   * @throws IllegalStateException
   *           when the accessToken parameter is not provided and no default
   *           accessToken was available.
   */
  private Token accessTokenFromArgs(Map args) throws IllegalStateException {
    Token accessToken = (Token) args.get("accessToken");
    if (accessToken == null) {
      accessToken = defaultAccessToken;
    }
    if (accessToken == null) {
      accessToken = new Token("", "");
    }
    return accessToken;
  }

  /**
   * Grabs the uri from the provided argument map, or returns the default URI.
   * 
   * @param args
   *          the arguments passed to either get or post
   * @return the provided uri
   * @throws IllegalStateException
   *           when the 'uri' parameters is not provided and no default uri was
   *           available.
   */
  private String uriParamFromArgs(Map args) throws IllegalStateException {
    String uriParam = (String) args.get("uri");
    
    if (uriParam == null) {
      uriParam = delegate.getUri().toString();
    }
    if (uriParam == null) {
      throw new IllegalStateException(
          "Default URI is null and no 'uri' parameter was provided"
      );
    }
    return uriParam;
  }

  /**
   * Adds the OAuth Authorization header to the wrapped HTTPBuilder.
   * 
   * @param authHeader
   *          the authorization header value
   * @param args
   *          the map of arguments passed to post/get
   */
  private void addAuthorizationHeader(String authHeader, Map args) {
    Map headers = (Map) args.get("headers");
    if (headers == null) {
      headers = new HashMap<String, String>();
      args.put("headers", headers);
    }
    headers.put("Authorization", authHeader);
  }

  /**
   * Constructs the authorization header from OAuth parameters.
   * 
   * @param method
   *          the request's HTTP method
   * @param uri
   *          the request's URI
   * @param accessToken
   *          the token for signing: can be access token or request token.
   * @param queryOrBodyParams
   *          the DECODED parameters for oauth to sign. With GET these are the
   *          query parameters. With POST both query parameters and body
   *          parameters (if this is a form-encoded request).
   * @return the signed OAuth "Authorization" header.
   */
  private String makeAuthHeader( Verb method, 
                                 String uri, 
                                 Token 
                                 accessToken,
                                 Map<String, ?> queryOrBodyParams ) {
    CustomRequest req = new CustomRequest(method, uri);
    for (Map.Entry param : queryOrBodyParams.entrySet()) {
      req.addBodyParameter( (String) param.getKey(), 
                            param.getValue().toString() );
    }
    signer.signRequest(req, accessToken);
    return req.authorization;
  }

  /**
   * Custom Scribe Request subclass gives us access to the Authorization header
   * when Scribe sets it.
   */
  protected static class CustomRequest extends Request {
    public CustomRequest(Verb verb, String url) {
      super(verb, url);
    }

    public String authorization = null;

    @Override
    public void addHeader(String key, String value) {
      if ("Authorization".equals(key)) {
        authorization = value;
      }
      super.addHeader(key, value);
    }
  }
}