import os
import requests
import pandas as pd
import io
from seleniumwire import webdriver
from selenium.webdriver.chrome.service import Service
#from selenium.webdriver.chrome.options import Options
from selenium.common.exceptions import TimeoutException
from selenium.common.exceptions import UnexpectedAlertPresentException
from selenium.common.exceptions import WebDriverException
import numpy as np

os.chdir(r'../Data/Data_website/Scrapping')

chrome_options = webdriver.ChromeOptions() 
#chrome_options.setPageLoadStrategy(PageLoadStrategy.NONE)
chrome_options.add_argument("--headless=chrome")
chrome_options.add_argument("start-maximized")
chrome_options.add_argument("enable-automation")
chrome_options.add_argument("--no-sandbox")
chrome_options.add_argument("--disable-dev-shm-usage")
chrome_options.add_argument("--disable-browser-side-navigation")
chrome_options.add_argument("--disable-gpu")
chrome_options.add_argument("--dns-prefetch-disable")

chrome_options.add_argument('allow-running-insecure-content')
chrome_options.add_argument('ignore-certificate-errors')
chrome_options.add_argument('allow-insecure-localhost')
chrome_options.add_argument('unsafely-treat-insecure-origin-as-secure')
chrome_options.add_argument('--ignore-ssl-errors=yes')
chrome_options.add_argument('--allow-insecure-localhost')
chrome_options.add_argument('--allow-running-insecure-content')
chrome_options.set_capability("acceptInsecureCerts", True)

user_agent = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36'
chrome_options.add_argument("user-agent="+user_agent)
PROXY = {
    "http": "xxx",
    "https": "xxx" #provide your proxy here
}
chrome_options.add_argument('--proxy-server=%s' %PROXY)
ser = Service(r'../Data/Data_website/Scrapping/chromedriver.exe')

def wayback_query(host, matchType, collapseType, proxy):
    query = 'http://web.archive.org/cdx/search/cdx?url={}&matchType={}&collapse={}&from=19960101&to=20221231&filter=mimetype:text/html&output=json'.format(host,matchType,collapseType)
    req = requests.get(query, headers={'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36'}, proxies=proxy)
    if req.status_code==200:
        df = pd.read_json(io.BytesIO(req.content))
        if df.shape[0]==0:
            return None
        else:
            df.columns = df.iloc[0]
            df = df.drop(df.index[0]).reset_index(drop=True)
            df['datetime'] = pd.to_datetime(df['timestamp'], format='%Y%m%d%H%M%S', errors='coerce')
            df['year'] = df['datetime'].dt.year
            if df['datetime'].isnull().any():
                df = df[df['datetime'].notnull()]
            df['baseKey'] = sorted(df['urlkey'].values, key=len)[0]
            df['url'] = ['https://web.archive.org/web' + '/' + t + '/' + o for t, o in zip(df['timestamp'].values, df['original'].values)]
            df.drop(columns=['timestamp','original','mimetype'], inplace=True)
            df.sort_values(['urlkey', 'datetime'], inplace=True)
            
            return df
    else:
        return None

def urlListgenerator(host, matchType, collapseType, proxy, chromeOptions, service):
    df = wayback_query(host, matchType, collapseType, proxy)
    if type(df) is type(None):
        return None
    else:
        minYear_alldata = int(df['year'].min())
        maxYear_alldata = int(df['year'].max())
        maxYear_baseKey = int(df[df.baseKey==df.urlkey]['year'].max())
        if maxYear_alldata<=maxYear_baseKey:
            df['baseKey4Yr'] = df.iloc[0]['baseKey']
            return df
        else:
            dfURL = pd.DataFrame({'year':pd.Series(range(minYear_alldata,maxYear_alldata+1))})
            dfURL['baseURL'] = host
            dfURL['baseKey4Yr'] = df.iloc[0]['baseKey']
            for y in range(maxYear_baseKey,maxYear_alldata+1):
                if len(df[df.year==y])==0:
                    dfURL.loc[dfURL.year==y, 'baseURL'] = 'NA'
                else:
                    lastDate = str(pd.to_datetime(df[df.year==y]['datetime'].max()).strftime('%Y%m%d%H%M%S'))
                    origWBURL = 'https://web.archive.org/web/'+lastDate+'/'+host
                    driver = webdriver.Chrome(options=chromeOptions, service=service, seleniumwire_options=proxy)
                    try: 
                        driver.get(origWBURL)
                        newURL = ":".join("/".join(driver.current_url.split("/")[:3]).split(":")[:2])
                        driver.close()
                        dfURL.loc[dfURL.year==y, 'baseURL'] = newURL
                    except (TimeoutException, UnexpectedAlertPresentException, WebDriverException):
                        driver.close()
                        dfURL.loc[dfURL.year==y, 'baseURL'] = 'NA'
            addnURLs = dfURL['baseURL'].unique()
            addnURLs = np.delete(addnURLs,np.where((addnURLs==host)|(addnURLs=='https://web.archive.org')))
            for elem in range(len(addnURLs)):
                newList = wayback_query(addnURLs[elem], matchType, collapseType, proxy)
                if type(newList) is type(None):
                    dfURL.loc[dfURL.baseURL==addnURLs[elem],'baseKey4Yr'] = "NA"
                else:
                    df = pd.concat([df,newList], axis=0, ignore_index=True)
                    dfURL.loc[dfURL.baseURL==addnURLs[elem],'baseKey4Yr'] = newList.iloc[0]['baseKey']
            dfURL.loc[dfURL.baseURL=='NA','baseKey4Yr'] = "NA"
            dfURL.drop('baseURL', axis=1, inplace=True)
            df = df.merge(dfURL, on='year')
            return df

#Reading in the files
#PE Firms
base_path = r'.\Inputs to the scraping process\rounds'

file_pattern = 'round {}.txt'
dataframes = []
for i in range(1, 8):
    file_path = os.path.join(base_path, file_pattern.format(i))
    if os.path.exists(file_path):
        df = pd.read_csv(file_path)
        dataframes.append(df)
    else:
        print(f"File not found: {file_path}")
websiteListPE = pd.concat(dataframes, ignore_index=True)

#HF Firms
websiteListHF = pd.read_csv(os.path.join(base_path, 'hedgeFunds.txt'))

#Used the below as an example
websiteListPE = pd.read_csv(r'.\Inputs to the scraping process\rounds\roundPE example.txt')
websiteListHF = pd.read_csv(r'.\Inputs to the scraping process\rounds\roundHF example.txt')

###PE Firms
for i in range(len(websiteListPE)):
    name = websiteListPE.iloc[i][0]
    if os.path.isfile('./URLlist/PE Firms/%s.pkl' %name)==True:
        continue
    else:
        print('PE Firm', i, websiteListPE.iloc[i][0], websiteListPE.iloc[i][1])
        urlList = urlListgenerator(websiteListPE.iloc[i][1], 'domain', 'digest', PROXY, chrome_options, ser)
        if type(urlList) is type(None):
            continue
        else:
            urlList['fm_id'] = websiteListPE.iloc[i][0]
            urlList.to_pickle('./URLlist/PE Firms/%s.pkl' %name, compression='zip')

###HF Firms
for i in range(len(websiteListHF)):
    name = websiteListHF.iloc[i][0]
    if os.path.isfile('./URLlist/HF Firms/%s.pkl' %name)==True:
        continue
    else:
        print('HF Firm', i, websiteListHF.iloc[i][0], websiteListHF.iloc[i][1])
        urlList = urlListgenerator(websiteListHF.iloc[i][1], 'domain', 'digest', PROXY, chrome_options, ser)
        if type(urlList) is type(None):
            continue
        else:
            urlList['fm_id'] = websiteListHF.iloc[i][0]
            urlList.to_pickle('./URLlist/HF Firms/%s.pkl' %name, compression='zip')