#!/usr/bin/env python3
import json
import os
import re
from pathlib import Path

import uvicorn
from fastapi import FastAPI, Body
from pysafebrowsing import SafeBrowsing
from starlette.responses import RedirectResponse, HTMLResponse, FileResponse, PlainTextResponse

app = FastAPI()


links: dict[str, str]
last_id: int

data_path = Path('data')
links_path = data_path / 'short-links.json'
last_id_path = data_path / 'last_id.txt'

# Removed lower l
chars = 'abcdefghijkmnopqrstuvwxyz'
base = len(chars)

# URL checks
re_url = re.compile(r"""^https?://(?:\S+(?::\S*)?@)?(?:(?!(?:10|127)(?:\.\d{1,3}){3})(?!(?:169\.254|192\.168)(?:\.\d{1,3}){2})(?!172\.(?:1[6-9]|2\d|3[0-1])(?:\.\d{1,3}){2})(?:[1-9]\d?|1\d\d|2[01]\d|22[0-3])(?:\.(?:1?\d{1,2}|2[0-4]\d|25[0-5])){2}(?:\.(?:[1-9]\d?|1\d\d|2[0-4]\d|25[0-4]))|(?:(?:[a-z\u00a1-\uffff0-9]-*)*[a-z\u00a1-\uffff0-9]+)(?:\.(?:[a-z\u00a1-\uffff0-9]-*)*[a-z\u00a1-\uffff0-9]+)*(?:\.(?:[a-z\u00a1-\uffff]{2,}))\.?)(?::\d{2,5})?(?:[/?#]\S*)?$""")
safe_browsing = SafeBrowsing(os.environ['GOOGLE_API_KEY'])


def store():
    data_path.mkdir(parents=True, exist_ok=True)
    links_path.write_text(json.dumps(links))
    last_id_path.write_text(str(last_id))


def load():
    global links
    global last_id
    links = json.loads(links_path.read_text()) if links_path.is_file() else {}
    last_id = int(last_id_path.read_text()) if last_id_path.is_file() else 1000


def encode(plain: int) -> str:
    cipher = ''
    while plain != 0:
        digit = int(plain % base)
        cipher += chars[digit]
        plain = (plain - digit) // base
    return cipher


def decode(s: str) -> int:
    total = 0
    for i, token in enumerate(s):
        digit = chars.index(token)
        assert digit != -1, f'Unknown char: {token}'
        total += digit * int(base ** i)
    return total


@app.get('/{short}')
def expand(short: str):
    print(short)
    print(links)
    if short in links:
        return RedirectResponse(links[short])
    else:
        return HTMLResponse('404', status_code=404)


@app.get('/')
def index():
    return FileResponse('index.html')


@app.get('/favicon.ico')
def favicon():
    return FileResponse('favicon.ico')


@app.put('/')
def put(name: str | None = None, body: str = Body()):
    try:
        global last_id

        # Check valid html
        assert re_url.match(body)
        sb = safe_browsing.lookup_url(body)
        print(sb)
        assert not sb['malicious'], f'Link is malicious ({",".join(sb["threats"]).lower()})'

        # Generate name
        while not name:
            last_id += 1
            name = encode(last_id)
            if name in links:
                name = None

        # Put name
        links[name] = body
        store()

        return PlainTextResponse(f'/{name}')

    except AssertionError as e:
        return PlainTextResponse(f'Error: {e}', status_code=400)

if __name__ == '__main__':
    load()
    uvicorn.run(app)