from serpapi import GoogleSearch
import pandas as pd
import numpy as np


def get_serp_result(keyword: str,
                    api_key: str,
                    n=10,
                    offset=0,
                    location: str = 'United States',
                    country: str = 'us',
                    language: str = 'en',
                    device: str = 'desktop'):
    """
    :param keyword: the search term you want info for
    :param api_key: your SerpApi API key
    :param n: the number of results you'd like to return
    :param offset: the offset number (i.e. an offset of 10 would give you results starting from page 2)
    :param location: a string of the location you want to search from
    :param country: the country code for the country you want to search from
    :param language: the language code for results desired
    :param device: the device type ('desktop', 'tabl' (for iPads), or 'mobile' (for iPhones))

    :return: a dictionary of the results from SerpApi
    """
    keyword_params = {
        "engine": "google",
        "q": keyword,
        "location": location,
        "google_domain": "google.com",
        "gl": country,
        "hl": language,
        "num": n,
        "start": offset,
        "device": device,
        "api_key": api_key
    }
    search = GoogleSearch(keyword_params)
    return search.get_dict()


def get_batched_organic_results(keyword_list, offsets):
    """
    :param keyword_list: a list of keywords you want to search for
    :param offsets: a list of offset values (for if you want to search multiple pages per keyword)
                    an offset list of [0, 10, 20] would give you pages 1, 2, and 3
    :return: a dataframe of all results
    """
    results = pd.DataFrame()

    for q in keyword_list:
        for o in offsets:
            temp_results = get_serp_result(keyword=q, n=10, offset=o)
            temp_df = extract_organic_df(keyword=q, result_set=temp_results)

            # Error handling for when no results are returned
            if temp_df is None:
                print(q + ', page ' + str(int(o / 10 + 1, )) + ': ERROR; no results returned')
                continue
            else:
                temp_df.insert(1, 'page', int(o / 10 + 1, ))
                print(q + ', page ' + str(int(o / 10 + 1, )) + ': ' + str(len(temp_df)) + ' results returned')
                results = pd.concat([results, temp_df], axis=0)

    return results


def extract_organic_df(result_set: pd.DataFrame):
    """
    :param result_set: a dictionary of SerpApi results generated by get_serp_result()
    :return:a dataframe of organic results, or None if there are none
    """
    # Error handling for when no results are returned
    try:
        organic = pd.DataFrame.from_dict(result_set['organic_results'])
        organic.insert(0, 'keyword', np.repeat(result_set['search_parameters']['q'], len(organic)))
        organic = organic[['keyword', 'position', 'link', 'title', 'snippet']]
        return organic
    except NameError:
        return None
    except KeyError:
        return None


def extract_knowledge_graph(result_set: pd.DataFrame):
    """
    :param result_set: a dictionary of SerpApi results generated by get_serp_result()
    :return:a dataframe of knowledge graph results, or None if there are none
    """
    # Error handling for when no results are returned
    try:
        out = result_set['knowledge_graph']
        out = pd.DataFrame({key: str(value) for key, value in out.items()}, index=[0])
        out.insert(0, 'keyword', np.repeat(result_set['search_parameters']['q'], len(out)))
        return out
    except NameError:
        return None
    except KeyError:
        return None


def extract_related_questions(result_set: pd.DataFrame):
    """
    :param result_set: a dictionary of SerpApi results generated by get_serp_result()
    :return:a dataframe of related questions results, or None if there are none
    """
    # Error handling for when no results are returned
    try:
        out = pd.DataFrame.from_dict(result_set['related_questions'])
        out.insert(0, 'keyword', np.repeat(result_set['search_parameters']['q'], len(out)))
        return out
    except NameError:
        return None
    except KeyError:
        return None


def extract_related_searches(result_set: pd.DataFrame):
    """
    :param result_set: a dictionary of SerpApi results generated by get_serp_result()
    :return:a dataframe of related searches results, or None if there are none
    """
    # Error handling for when no results are returned
    try:
        out = pd.DataFrame.from_dict(result_set['related_searches'])
        out.insert(0, 'keyword', np.repeat(result_set['search_parameters']['q'], len(out)))
        return out
    except NameError:
        return None
    except KeyError:
        return None


def extract_top_stories(result_set: pd.DataFrame):
    """
    :param result_set: a dictionary of SerpApi results generated by get_serp_result()
    :return:a dataframe of top stories results, or None if there are none
    """
    # Error handling for when no results are returned
    try:
        out = pd.DataFrame.from_dict(result_set['top_stories'])
        out.insert(0, 'keyword', np.repeat(result_set['search_parameters']['q'], len(out)))
        return out
    except NameError:
        return None
    except KeyError:
        return None
