import asyncio
import re
import sys
from typing import Union

import aiohttp
from bs4 import BeautifulSoup

from .helpers import SB, random_string
from .types import Language


class CodeforcesClient:
    def __init__(self, cookies: dict[str, str], csrf_token: str, ftaa: str, user_agent: str):
        self.cookies = cookies
        self.user_agent = user_agent

        self._csrf_token = csrf_token
        self._ftaa = ftaa


    async def __aenter__(self):
        self._session = aiohttp.ClientSession()
        self._session.cookie_jar.update_cookies(self.cookies)
        self._session.headers['User-Agent'] = self.user_agent
        return self


    async def __aexit__(self, *_):
        await self.close()


    @staticmethod
    def auth(username: str, password: str): # TODO: error checking
        with SB() as sb:
            sb.uc_open_with_reconnect("https://codeforces.com/enter", 2)
            sb.uc_gui_click_captcha()

            sb.click('a[href="/enter?back=%2F"]')

            sb.type("#handleOrEmail", username)
            sb.type("#password", password)

            sb.click("#remember")
            sb.click('input[type="submit"]')

            sb.wait_for_element("#jGrowl")

            cookies = sb.get_cookies()
            user_agent = sb.get_user_agent()
            source = sb.get_page_source()

        match = re.findall(r'<meta name="X-Csrf-Token" content="([a-f0-9]+)">', source)
        assert match
        csrf = match[0]

        match = re.findall(r'window._ftaa = "(\w+)";', source)
        assert match
        ftaa = match[0]

        cookies = {cookie['name']: cookie['value'] for cookie in cookies}

        return CodeforcesClient(cookies, csrf, ftaa, user_agent)


    @staticmethod
    def find_error(body: str) -> Union[str, None]:
        r = re.compile(r'error[a-zA-Z_\-\ ]*">(.*?)</span>')
        m = re.findall(r, body)
        return m[0] if m else None


    async def submit(self, code: str, contest: int, task: str, lang: Language): # TODO: error checking
        r = await self._session.post(f"https://codeforces.com/contest/{contest}/submit?csrf_token={self._csrf_token}", data = {
            "csrf_token": self._csrf_token,
            "ftaa": self._ftaa,
            "bfaa": random_string(32, '[a-f0-9]'),
            "action": "submitSolutionFormSubmitted",
            "submittedProblemIndex": task, # TODO: idx saving
            "programTypeId": lang.value,
            "source": code,
            "tabSize": "4",
            "sourceFile": "",
            "_tta": "239",
        })

        err = self.find_error(await r.text())

        if err:
            print(err, file=sys.stderr)


    async def parse_tests(self, contest: int, tasks: list[str]):
        futures = [self._session.get(f"https://codeforces.com/contest/{contest}/problem/{task}") for task in tasks]
        requests = await asyncio.gather(*futures)

        for idx, r in enumerate(requests):
            soup = BeautifulSoup(await r.text(), 'html.parser')
            test = soup.select_one('div.sample-test')
            assert test

            # funny line haha
            inputs = ['\n'.join([el.text for el in input.select('div.test-example-line')] or [input.text]).strip() for input in test.select('div.input > pre')]
            outputs = ['\n'.join([el.text for el in output.select('div.test-example-line')] or [output.text]).strip() for output in test.select('div.output > pre')]

            yield {tasks[idx].lower(): [*zip(inputs, outputs)]}


    async def parse_tasks(self, contest: int):
        r = await self._session.get(f"https://codeforces.com/contest/{contest}")
        soup = BeautifulSoup(await r.text(), 'html.parser')

        table = soup.select_one('table.problems')
        assert table

        elements = table.select('tr')[1:]

        out = []

        for item in elements:
            status = None
            if 'class' in item.attrs:
                status = item.attrs['class'][0]
            inner_items = item.select('td')
            assert inner_items
            name_and_constraints = inner_items[1].select('div > div')
            id = inner_items[0].text.strip()

            out.append({
                'id': id,
                'name': name_and_constraints[0].text.strip(),
                'status': status,
                'constraints': [*name_and_constraints[1]][2].text.strip(),
            })

        return out


    def to_dict(self):
        return {
            'ftaa': self._ftaa,
            'cookies': self.cookies,
            'user_agent': self.user_agent,
            'csrf_token': self._csrf_token
        }


    async def close(self):
        await self._session.close()