#!/usr/bin/env python3
import base64
import hashlib
import json
import sys
import xml.etree.ElementTree as ET
from urllib import parse

sys.path.insert(0, "/usr/local/mgr5/lib/python")
import billmgr.exception as exc
import billmgr.logger as logging
import requests
from billmgr.modules.processing import Module

logging.init_logging("omvkid")
logger = logging.get_logger("omvkid")

REQUIRED_PARAMS = ['device_id']

CLIENT_ID = None
SERVICE_TOKEN = ""


def load_config():
    global CLIENT_ID, SERVICE_TOKEN

    if CLIENT_ID is not None:
        return

    with open("/usr/local/mgr5/etc/oauth/vkid.json") as f:
        raw = json.load(f)

    CLIENT_ID = raw["client_id"]
    SERVICE_TOKEN = raw.get("service_token", raw.get("secret", ""))

mod = Module()
mod.add_argument("--state", type=str, help="state", dest="state")
mod.add_argument("--code", type=str, help="auth code from VK", dest="code")
mod.add_argument("--response_type", type=str, help="response type", dest="response_type")
mod.add_argument("--return_url", type=str, help="url for return after authorize", dest="return_url")
mod.add_argument("--from", type=str, help="from", dest="from")
mod.add_argument("--ip", type=str, help="ip", dest="ip")
mod.add_argument("--params", type=str, help="extra session params as JSON", dest="params", default="{}")


def derive_pkce(state: str):
    """Deterministically computes code_verifier/code_challenge from state and SERVICE_TOKEN, 
    eliminating the need to persist the verifier between make_url and get_user_data."""
    digest = hashlib.sha256((SERVICE_TOKEN + state).encode()).digest()
    verifier = base64.urlsafe_b64encode(digest).decode().rstrip("=")
    challenge_digest = hashlib.sha256(verifier.encode()).digest()
    challenge = base64.urlsafe_b64encode(challenge_digest).decode().rstrip("=")
    return verifier, challenge


@mod.command("make_url")
def make_url(state: str, response_type: str, return_url: str):
    load_config()
    _, code_challenge = derive_pkce(state)

    params = {
        "client_id": CLIENT_ID,
        "scope": "vkid.personal_info email",
        "response_type": response_type,
        "redirect_uri": return_url,
        "state": state,
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
    }
    url = f"https://id.vk.ru/authorize?{parse.urlencode(params)}"
    logger.info("URL: '%s'", url)

    doc = ET.Element("doc")
    ET.SubElement(doc, "url").text = url
    ET.dump(doc)


@mod.command("get_user_data")
def get_user_data(code: str, return_url: str, state: str, params: str = "{}"):
    load_config()
    session_params = json.loads(params or "{}")
    device_id = session_params.get("device_id", "")

    if not device_id:
        logger.error("device_id is missing in session params: %s", session_params)
        raise exc.XmlException("token_exchange", "missed", "device_id")

    logger.info("get_user_data started for state %s, device_id=%s", state, device_id)
    code_verifier, _ = derive_pkce(state)

    logger.info("Exchange code for tokens")
    r = requests.post(
        "https://id.vk.ru/oauth2/auth",
        data={
            "grant_type": "authorization_code",
            "client_id": CLIENT_ID,
            "service_token": SERVICE_TOKEN,
            "code": code,
            "code_verifier": code_verifier,
            "redirect_uri": return_url,
            "device_id": device_id,
        },
    )
    logger.info("Token response status: %d, body: %s", r.status_code, r.text)

    token_data = r.json()
    if "error" in token_data:
        logger.error("Token exchange error: %s", token_data)
        raise exc.XmlException(
            "token_exchange", token_data["error"], token_data.get("error_description", "")
        )

    access_token = token_data.get("access_token")
    if not access_token:
        logger.error("No access_token in response")
        raise exc.XmlException("token_exchange", "missed", "access_token")

    logger.info("Fetch user info")
    r = requests.post(
        "https://id.vk.ru/oauth2/user_info",
        data={
            "access_token": access_token,
            "client_id": CLIENT_ID,
        },
    )
    logger.info("User info response: %s", r.text)

    user_info = r.json()
    if "error" in user_info:
        logger.error("User info error: %s", user_info)
        raise exc.XmlException(
            "user_info", user_info["error"], user_info.get("error_description", "")
        )

    user = user_info.get("user")
    if not user:
        logger.error("No 'user' field in user_info response")
        raise exc.XmlException("user_info", "missed", "user")

    user_id = str(user.get("user_id", ""))
    email = user.get("email", "")
    first_name = user.get("first_name", "")
    last_name = user.get("last_name", "")
    phone = user.get("phone", "")

    if not email:
        logger.info("Email not found in user_info, trying VK API users.get")
        r = requests.get(
            "https://api.vk.ru/method/users.get",
            params={
                "v": "5.131",
                "access_token": access_token,
                "fields": "email",
            },
        )
        logger.info("VK API users.get response: %s", r.text)
        vk_response = r.json()
        if "error" not in vk_response and "response" in vk_response:
            vk_user = vk_response["response"][0]
            email = vk_user.get("email", "")
            if not first_name:
                first_name = vk_user.get("first_name", "")
            if not last_name:
                last_name = vk_user.get("last_name", "")
            if not user_id:
                user_id = str(vk_user.get("id", ""))
        else:
            logger.warning("VK API users.get failed: %s", vk_response.get("error", "unknown error"))

    if not user_id or not email or not first_name:
        logger.error(
            "Missing required fields: user_id=%s, email=%s, first_name=%s",
            user_id,
            email,
            first_name,
        )
        raise exc.XmlException("user_info", "missed", "required_fields")

    doc = ET.Element("doc")
    info = ET.SubElement(doc, "userdata")
    ET.SubElement(info, "id").text = user_id
    ET.SubElement(info, "email").text = email
    ET.SubElement(info, "realname").text = f"{first_name} {last_name}"
    ET.SubElement(info, "firstname").text = first_name
    ET.SubElement(info, "lastname").text = last_name
    if phone:
        ET.SubElement(info, "phone").text = phone
    ET.dump(doc)
    logger.info("User data successfully returned")


if __name__ == "__main__":
    mod.run()