#!/usr/bin/env python3
import json
import os
import time
import uuid
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib.parse import parse_qs, urlencode, urlparse

import jwt
from cryptography.hazmat.primitives.asymmetric import rsa


PORT = int(os.environ.get("CASAN_IDP_PORT", "8080"))
ISSUER = os.environ.get("CASAN_IDP_ISSUER", f"http://127.0.0.1:{PORT}")
DEFAULT_SUB = os.environ.get("CASAN_IDP_SUB", "oidc-ops")
DEFAULT_EMAIL = os.environ.get("CASAN_IDP_EMAIL", "oidc-ops@example.com")
DEFAULT_GROUPS = [g for g in os.environ.get("CASAN_IDP_GROUPS", "casan-org-admin,casan-approver").split(",") if g]
KID = "casan-local-prod-idp"
KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048)
CODES = {}


def b64u_int(value: int) -> str:
    raw = value.to_bytes((value.bit_length() + 7) // 8, "big")
    import base64

    return base64.urlsafe_b64encode(raw).decode().rstrip("=")


def jwk():
    numbers = KEY.public_key().public_numbers()
    return {
        "kty": "RSA",
        "use": "sig",
        "alg": "RS256",
        "kid": KID,
        "n": b64u_int(numbers.n),
        "e": b64u_int(numbers.e),
    }


def sign_claims(client_id="casan-control-panel", nonce="", ttl_s=300):
    now = int(time.time())
    claims = {
        "iss": ISSUER,
        "sub": DEFAULT_SUB,
        "aud": client_id,
        "email": DEFAULT_EMAIL,
        "groups": DEFAULT_GROUPS,
        "iat": now,
        "exp": now + int(ttl_s),
    }
    if nonce:
        claims["nonce"] = nonce
    token = jwt.encode(claims, KEY, algorithm="RS256", headers={"kid": KID})
    return token, claims


class Handler(BaseHTTPRequestHandler):
    def send_json(self, code, payload):
        body = json.dumps(payload).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):  # noqa: N802
        parsed = urlparse(self.path)
        path = parsed.path
        if path == "/healthz":
            self.send_json(200, {"status": "ok", "issuer": ISSUER})
        elif path == "/.well-known/openid-configuration":
            self.send_json(
                200,
                {
                    "issuer": ISSUER,
                    "authorization_endpoint": f"{ISSUER}/authorize",
                    "jwks_uri": f"{ISSUER}/.well-known/jwks.json",
                    "token_endpoint": f"{ISSUER}/token",
                    "userinfo_endpoint": f"{ISSUER}/userinfo",
                    "id_token_signing_alg_values_supported": ["RS256"],
                    "response_types_supported": ["code"],
                    "subject_types_supported": ["public"],
                    "scopes_supported": ["openid", "email", "profile", "groups"],
                    "claims_supported": ["sub", "email", "groups"],
                },
            )
        elif path == "/.well-known/jwks.json":
            self.send_json(200, {"keys": [jwk()]})
        elif path == "/authorize":
            q = parse_qs(parsed.query)
            redirect_uri = q.get("redirect_uri", [""])[0]
            if not redirect_uri:
                self.send_json(400, {"error": "missing_redirect_uri"})
                return
            code = uuid.uuid4().hex
            CODES[code] = {
                "client_id": q.get("client_id", ["casan-control-panel"])[0],
                "nonce": q.get("nonce", [""])[0],
                "sub": DEFAULT_SUB,
                "email": DEFAULT_EMAIL,
                "groups": DEFAULT_GROUPS,
            }
            params = {"code": code}
            if q.get("state"):
                params["state"] = q["state"][0]
            sep = "&" if "?" in redirect_uri else "?"
            self.send_response(302)
            self.send_header("Location", redirect_uri + sep + urlencode(params))
            self.end_headers()
        elif path == "/userinfo":
            self.send_json(200, {"sub": DEFAULT_SUB, "email": DEFAULT_EMAIL, "groups": DEFAULT_GROUPS})
        else:
            self.send_json(404, {"error": "not_found"})

    def do_POST(self):  # noqa: N802
        if urlparse(self.path).path != "/token":
            self.send_json(404, {"error": "not_found"})
            return
        size = int(self.headers.get("Content-Length", "0"))
        raw = self.rfile.read(size) or b"{}"
        ctype = self.headers.get("Content-Type", "")
        if "application/json" in ctype:
            try:
                payload = json.loads(raw)
            except json.JSONDecodeError:
                self.send_json(400, {"error": "invalid_json"})
                return
            # Backward-compatible approval-verifier token minting.
            now = int(time.time())
            claims = {
                "iss": ISSUER,
                "sub": payload.get("sub", DEFAULT_SUB),
                "role": payload.get("role", "ops"),
                "action": payload.get("action", "deploy"),
                "actor": payload.get("actor", "alice"),
                "input_sha256": payload.get("input_sha256", ""),
                "iat": now,
                "exp": now + int(payload.get("ttl_s", 300)),
            }
            token = jwt.encode(claims, KEY, algorithm="RS256", headers={"kid": KID})
            self.send_json(200, {"access_token": token, "token_type": "Bearer", "expires_in": claims["exp"] - now})
            return

        form = parse_qs(raw.decode())
        code = form.get("code", [""])[0]
        rec = CODES.pop(code, None)
        if not rec:
            self.send_json(400, {"error": "invalid_grant"})
            return
        token, claims = sign_claims(client_id=rec.get("client_id") or form.get("client_id", ["casan-control-panel"])[0],
                                   nonce=rec.get("nonce", ""))
        self.send_json(200, {
            "access_token": token,
            "id_token": token,
            "token_type": "Bearer",
            "expires_in": claims["exp"] - int(time.time()),
        })

    def log_message(self, *args):
        pass


if __name__ == "__main__":
    HTTPServer(("0.0.0.0", PORT), Handler).serve_forever()