import json
import uuid
import requests
import time
import secrets

from pathlib import Path
from http.server import BaseHTTPRequestHandler, HTTPServer


def get_copilot_token_from_neovim_vault():
    """Read Copilot token from Neovim's vault_state.json."""
    vault_path = Path.home() / ".local/share/lazyvim/gp/persisted/vault_state.json"

    try:
        with open(vault_path, 'r') as f:
            vault = json.load(f)
            return vault['copilot_bearer']['token']
    except (FileNotFoundError, KeyError, json.JSONDecodeError) as e:
        raise Exception(f"Failed to read Copilot token from vault: {e}")


def gen_hex_str(n):
    return secrets.token_hex(n)


class CopilotOpenAIProxy(BaseHTTPRequestHandler):
    DEFAULT_HEADERS = {
        'Editor-Version': "vscode/1.83.1",
        'Editor-Plugin-Version': "copilot-chat/0.8.0",
        'Openai-Organization': "github-copilot",
        'Openai-Intent': "conversation-panel",
        'Content-Type': "text/event-stream; charset=utf-8",
        'User-Agent': "GitHubCopilotChat/0.8.0",
        'Accept': "*/*",
        'Accept-Encoding': "gzip,deflate,br",
    }
    copilot_token = None

    def __init__(self, token):
        self.copilot_token = token

    def __call__(self, *args, **kwargs):
        """Handle a request."""
        super().__init__(*args, **kwargs)

    def request_id(self):
        return gen_hex_str(4) + "-" + gen_hex_str(2) + "-" + gen_hex_str(2) + "-" + gen_hex_str(2) + "-" + gen_hex_str(6)

    def session_id(self):
        return gen_hex_str(4) + "-" + gen_hex_str(2) + "-" + gen_hex_str(2) + "-" + gen_hex_str(2) + "-" + gen_hex_str(12)

    def machine_id(self):
        return gen_hex_str(32)

    def do_POST(self):
        if self.path != '/v1/chat/completions':
            self.send_response(
                requests.codes.method_not_allowed, "Method Not Allowed")
            self.end_headers()

            return

        content_length = int(self.headers['Content-Length'])
        body = self.rfile.read(content_length)
        copilot_resp = self.copilot_request(body)

        if copilot_resp.status_code != requests.codes.ok:
            self.send_response(copilot_resp.status_code)
            self.end_headers()

            return

        self.send_response(requests.codes.ok)
        self.end_headers()
        self.wfile.write(copilot_resp.content)
        print(copilot_resp.content)

    def copilot_request(self, body):
        url = "https://api.githubcopilot.com/chat/completions"

        headers = self.DEFAULT_HEADERS.copy()
        headers['Authorization'] = f'Bearer {self.copilot_token}'
        headers['X-Request-ID'] = self.request_id()
        headers['Vscode-Sessionid'] = self.session_id()
        headers['Vscode-Machineid'] = self.machine_id()

        return requests.post(url, headers=headers, data=body)
    pass


if __name__ == '__main__':
    token = get_copilot_token_from_neovim_vault()
    proxy = CopilotOpenAIProxy(token)

    server_address = ('127.0.0.1', 3040)
    httpd = HTTPServer(server_address, proxy)
    print("Proxy running at http://127.0.0.1:3040")
    httpd.serve_forever()