#!/usr/bin/python3
import urllib.request
import json

headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36'
}

class MyServer:
    def __init__(self, environ, start_response):
        self.environ = environ
        self.start_response = start_response
        self.mapping_file = 'mapping.json'
        self.mappings = self.load_mappings()
        self.default_hs = "https://matrix-client.matrix.org"
        self.blacklist = set(line.strip() for line in open('blacklist.txt'))

    def __iter__(self):

        path_info = hs = self.environ['PATH_INFO'].split('/')

        if path_info[2] == 'client':
            mediatype = path_info[5]
            hs = path_info[6]
            mediaid = path_info[7]
        else:
            mediatype = path_info[4]
            hs = path_info[5]
            mediaid = path_info[6]

        if not hs in self.blacklist:
            hsu = self.mappings.get(hs)

            if not hsu:
                wellknown = "https://" + hs + "/.well-known/matrix/client"
                req = urllib.request.Request(wellknown, headers=headers)
                try:
                    with urllib.request.urlopen(req, timeout=2) as cont:
                        hsu = json.load(cont)
                        hsu = hsu["m.homeserver"]["base_url"].rstrip('/')
                        self.mappings[hs] = hsu
                        self.save_mappings()
                except Exception as e: # I don't care to fix this properly, not my problem
                    hsu = self.default_hs
        else:
            hsu = self.default_hs

        hsp = hsu + '/_matrix/media/v3/' + mediatype + '/' + hs + '/' + mediaid + '?' + self.environ['QUERY_STRING']
        self.start_response('301 Moved Permanently', [('Location', hsp)])
        return iter([])

    def load_mappings(self):
        try:
            with open(self.mapping_file, 'r') as f:
                return json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            return {}

    def save_mappings(self):
        with open(self.mapping_file, 'w') as f:
            json.dump(self.mappings, f)

if __name__ == "__main__":
    from gunicorn.app.base import BaseApplication

    class GunicornServer(BaseApplication):
        def __init__(self, app, options=None):
            self.options = options or {}
            self.application = app
            super().__init__()

        def load_config(self):
            for key, value in self.options.items():
                self.cfg.set(key, value)

        def load(self):
            return self.application

    options = {
        'bind': 'localhost:9999',
        'workers': 48,  # Adjust the number of workers based on your system's resources - ChatGPT
    }

    server = GunicornServer(MyServer, options)
    server.run()