import asyncio
import json
import datetime
import threading
import time

import requests
from selenium import webdriver
from selenium.common.exceptions import JavascriptException
from selenium.webdriver import ActionChains
import chromedriver_binary
import websockets


ENVIRONMENTS = {
    'local': 'LOCAL',
    'cloud': 'DC',
    'unblockable': 'RESID',
}
FORMATS = {
    'json': 'json_format',
    'csv': 'csv_format',
}
OUTPUTS = set(['data', 'file'])
STARTING_LIMIT = 64
MAX_RESPONSE_THRESHOLD = 2 ** 20


class Client:
    def __init__(self, username, password, host='parsagon.io'):
        data = {'username': username, 'password': password}
        r = requests.post(f'https://{host}/api/accounts/token-auth/', json=data)
        if not r.ok:
            self._display_errors(r)
        self.token = r.json()['token']
        self.host = host

    def _display_errors(self, response):
        errors = response.json()
        if 'non_field_errors' in errors:
            raise Exception(errors['non_field_errors'])
        else:
            raise Exception(errors)

    async def handle_driver(self, result_id, exit_event):
        driver = webdriver.Chrome()
        async with websockets.connect(f'wss://{self.host}/ws/scrapers/results/{result_id}/client/') as websocket:
            while not exit_event.is_set():
                try:
                    message_str = await asyncio.wait_for(websocket.recv(), timeout=2)
                except asyncio.TimeoutError:
                    continue
                message = json.loads(message_str)
                convo_id = message['convo_id']
                response = 'OK'
                command = message['command']
                if command == 'get':
                    driver.get(message['url'])
                    await asyncio.sleep(2)
                elif command == 'mark':
                    elem_idx = driver.execute_script(f"let elemIdx = {message['elem_idx']}; for (const node of document.querySelectorAll(':not([data-parsagon-io-marked])')) {{ node.setAttribute('data-parsagon-io-marked', elemIdx); elemIdx++; }} return elemIdx;")
                    await websocket.send(json.dumps({'response': elem_idx, 'convo_id': convo_id}))
                    continue
                elif command == 'scroll':
                    driver.execute_script("window.scrollTo({top: document.body.scrollHeight, behavior: 'smooth'});")
                    await asyncio.sleep(1)
                elif command == 'click':
                    actions = ActionChains(driver)
                    target = driver.find_element_by_xpath(f"//*[@data-parsagon-io-marked={message['target_id']}]")
                    driver.execute_script("arguments[0].scrollIntoView({behavior: 'smooth', block: 'center', inline: 'center'});", target)
                    await asyncio.sleep(0.5)
                    try:
                        actions.move_to_element(target).click().perform()
                        await asyncio.sleep(2)
                        driver.execute_script("window.scrollTo({top: document.body.scrollHeight, behavior: 'smooth'});")
                        await asyncio.sleep(1)
                    except JavascriptException:
                        pass
                elif command == 'inspect':
                    query = message['query']
                    if query == 'url':
                        response = driver.current_url
                    elif query == 'page_source':
                        response = driver.page_source
                    elif query == 'target_data':
                        target = driver.find_element_by_xpath(f"//*[@data-parsagon-io-marked={message['target_id']}]")
                        tag = target.tag_name
                        text = target.text
                        href = target.get_attribute('href')
                        url = driver.execute_script(f"return document.querySelector('[data-parsagon-io-marked=\"{message['target_id']}\"]').href;")
                        response = {'tag': tag, 'text': text, 'href': href, 'url': url}
                await websocket.send(json.dumps({'response': response, 'convo_id': convo_id}))
        driver.quit()

    def run_driver(self, result_id, exit_event):
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(self.handle_driver(result_id, exit_event))
        loop.close()

    def get_result(self, result_id, format, output, file_path):
        headers = {'Content-Type': 'application/json', 'Authorization': f'Token {self.token}'}
        r = requests.post(f'https://{self.host}/api/scrapers/results/{result_id}/execute/', headers=headers)
        if not r.ok:
            self._display_errors(r)
        while True:
            r = requests.get(f'https://{self.host}/api/scrapers/results/{result_id}/', headers=headers)
            if not r.ok:
                self._display_errors(r)
            result_data = r.json()
            if result_data['status'] == 'FINISHED':
                break
            elif result_data['status'] == 'ERROR':
                raise Exception('A server error occurred. Please notify Parsagon.')
            time.sleep(5)

        if output == 'file':
            with open(file_path, 'w') as f:
                if format == 'csv':
                    raise Exception("Output type 'file' not yet supported for format 'csv'")
                else:
                    offset = 0
                    limit = STARTING_LIMIT
                    offset_incr = 0
                    max_response_size = 0
                    r = requests.get(f'https://{self.host}/api/scrapers/results/{result_id}/download/?data_format={format}&offset={offset}&limit={limit}', headers=headers)
                    if not r.ok:
                        self._display_errors(r)
                    result_data = r.json()
                    offset_incr = len(result_data['result'])
                    offset += offset_incr
                    new_data = json.dumps(result_data)
                    f.write(new_data[:-2])
                    max_response_size = max(max_response_size, len(new_data))

                    while offset_incr == limit:
                        if max_response_size < MAX_RESPONSE_THRESHOLD:
                            limit *= 2

                        r = requests.get(f'https://{self.host}/api/scrapers/results/{result_id}/download/?data_format={format}&offset={offset}&limit={limit}', headers=headers)
                        if not r.ok:
                            self._display_errors(r)
                        result_data = r.json()
                        offset_incr = len(result_data['result'])
                        if not offset_incr:
                            break
                        offset += offset_incr
                        new_data = json.dumps(result_data['result'])
                        new_data = ',' + new_data[1: -1]
                        f.write(new_data)
                        max_response_size = max(max_response_size, len(new_data))
                    f.write('}}')
        else:
            r = requests.get(f'https://{self.host}/api/scrapers/results/{result_id}/download/?data_format={format}',
                             headers=headers)
            if not r.ok:
                self._display_errors(r)
            data = r.json()
            if format == 'csv':
                return data['result']
            else:
                return data

    def execute(self, scraper_name, urls, env, max_page_loads=1, format='json', output='data', file_path=''):
        if env not in ENVIRONMENTS:
            raise ValueError("Environment must be 'local', 'cloud', or 'unblockable'")
        if format not in FORMATS:
            raise ValueError("Format must be 'json' or 'csv'")
        if output not in OUTPUTS:
            raise ValueError("Output must be 'data' or 'file'")
        if output == 'file' and not file_path:
            raise ValueError("Output type is 'file' but no file path was given")

        headers = {'Content-Type': 'application/json', 'Authorization': f'Token {self.token}'}

        data = {'scraper_name': scraper_name, 'urls': urls, 'max_page_loads': max_page_loads}
        r = requests.post(f'https://{self.host}/api/scrapers/runs/', headers=headers, json=data)
        if not r.ok:
            self._display_errors(r)
        run = r.json()

        if not run['scraper'][FORMATS[format]]:
            raise Exception(f'{format} format is unavailable for this scraper')

        data = {'environment': ENVIRONMENTS[env]}
        r = requests.post(f'https://{self.host}/api/scrapers/runs/{run["id2"]}/results/', headers=headers, json=data)
        if not r.ok:
            self._display_errors(r)
        result = r.json()

        exit_event = threading.Event()
        if env == 'local':
            driver_task = threading.Thread(target=self.run_driver, args=[result['id2'], exit_event])
            driver_task.start()
        return_value = self.get_result(result['id2'], format, output, file_path)
        exit_event.set()
        time.sleep(2.1)
        return return_value
