import type { NextFunction, Request, Response } from "express";
import { createHash, randomBytes } from "node:crypto";
import { createRemoteJWKSet, jwtVerify, SignJWT } from "jose";

let mcpJwks: ReturnType<typeof createRemoteJWKSet> | null = null;
let webJwks: ReturnType<typeof createRemoteJWKSet> | null = null;

const SESSION_COOKIE = "media_session";
const STATE_COOKIE = "media_oidc_state";
const VERIFIER_COOKIE = "media_oidc_verifier";
const NONCE_COOKIE = "media_oidc_nonce";
const RETURN_COOKIE = "media_return_to";

function publicUrl(req: Request) {
  return (process.env.PUBLIC_URL || `${req.protocol}://${req.get("host")}`).replace(/\/$/, "");
}

function base64url(value: Buffer) {
  return value.toString("base64url");
}

function randomToken(bytes = 32) {
  return base64url(randomBytes(bytes));
}

function readCookie(req: Request, name: string): string | null {
  const raw = req.get("cookie") ?? "";
  for (const part of raw.split(";")) {
    const [key, ...rest] = part.trim().split("=");
    if (key === name) return decodeURIComponent(rest.join("="));
  }
  return null;
}

function cookie(name: string, value: string, maxAgeSeconds: number) {
  return [
    `${name}=${encodeURIComponent(value)}`,
    "Path=/",
    "HttpOnly",
    "Secure",
    "SameSite=Lax",
    `Max-Age=${maxAgeSeconds}`
  ].join("; ");
}

function clearCookie(name: string) {
  return [
    `${name}=`,
    "Path=/",
    "HttpOnly",
    "Secure",
    "SameSite=Lax",
    "Max-Age=0"
  ].join("; ");
}

function safeReturnTo(value: unknown) {
  const text = String(value ?? "/");
  return text.startsWith("/") && !text.startsWith("//") ? text : "/";
}

function webConfig(req: Request) {
  const issuer = process.env.WEB_OIDC_ISSUER || process.env.KEYCLOAK_ISSUER;
  const clientId = process.env.WEB_OIDC_CLIENT_ID;
  const clientSecret = process.env.WEB_OIDC_CLIENT_SECRET;
  if (!issuer || !clientId || !clientSecret) {
    throw new Error("Web OIDC authentication is not configured");
  }
  return {
    issuer: issuer.replace(/\/$/, ""),
    clientId,
    clientSecret,
    redirectUri: `${publicUrl(req)}/auth/callback`
  };
}

function sessionKey() {
  const secret = process.env.WEB_OIDC_CLIENT_SECRET;
  if (!secret) throw new Error("Web OIDC client secret is unavailable");
  return createHash("sha256")
    .update(`media-intelligence-session:${secret}`)
    .digest();
}

async function createSession(claims: Record<string, unknown>) {
  return new SignJWT(claims)
    .setProtectedHeader({ alg: "HS256", typ: "JWT" })
    .setIssuer("media-intelligence-web")
    .setAudience("media-intelligence-web")
    .setIssuedAt()
    .setExpirationTime("8h")
    .sign(sessionKey());
}

async function verifySession(token: string) {
  const { payload } = await jwtVerify(token, sessionKey(), {
    issuer: "media-intelligence-web",
    audience: "media-intelligence-web"
  });
  return payload;
}

export function protectedResourceMetadata(req: Request) {
  const base = publicUrl(req);
  return {
    resource: `${base}/api/mcp`,
    authorization_servers: [
      process.env.KEYCLOAK_ISSUER || "https://auth.arrowstack.dev/realms/dev"
    ],
    bearer_methods_supported: ["header"],
    scopes_supported: ["openid", "profile", "email"]
  };
}

export async function requireMcpAuth(req: Request, res: Response, next: NextFunction) {
  if (process.env.AUTH_DISABLED === "true") return next();

  const issuer = process.env.KEYCLOAK_ISSUER;
  const clientId = process.env.KEYCLOAK_CLIENT_ID;
  if (!issuer || !clientId) {
    return res.status(503).json({ error: "MCP OAuth is not configured" });
  }

  const header = req.get("authorization") ?? "";
  if (!header.startsWith("Bearer ")) {
    res.setHeader(
      "WWW-Authenticate",
      `Bearer resource_metadata="${publicUrl(req)}/.well-known/oauth-protected-resource/api/mcp"`
    );
    return res.status(401).json({ error: "OAuth access token required" });
  }

  try {
    if (!mcpJwks) {
      mcpJwks = createRemoteJWKSet(
        new URL(`${issuer.replace(/\/$/, "")}/protocol/openid-connect/certs`)
      );
    }
    const { payload } = await jwtVerify(header.slice(7), mcpJwks, { issuer });
    const aud = Array.isArray(payload.aud) ? payload.aud : payload.aud ? [payload.aud] : [];
    if (payload.azp !== clientId && !aud.includes(clientId)) {
      return res.status(403).json({ error: "OAuth token was not issued for this MCP client" });
    }
    (req as any).auth = payload;
    return next();
  } catch {
    return res.status(401).json({ error: "Invalid OAuth access token" });
  }
}

export async function beginWebLogin(req: Request, res: Response) {
  try {
    const cfg = webConfig(req);
    const state = randomToken();
    const verifier = randomToken(48);
    const nonce = randomToken();
    const challenge = base64url(createHash("sha256").update(verifier).digest());
    const returnTo = safeReturnTo(req.query.returnTo);

    res.append("Set-Cookie", cookie(STATE_COOKIE, state, 600));
    res.append("Set-Cookie", cookie(VERIFIER_COOKIE, verifier, 600));
    res.append("Set-Cookie", cookie(NONCE_COOKIE, nonce, 600));
    res.append("Set-Cookie", cookie(RETURN_COOKIE, returnTo, 600));

    const url = new URL(`${cfg.issuer}/protocol/openid-connect/auth`);
    url.searchParams.set("client_id", cfg.clientId);
    url.searchParams.set("redirect_uri", cfg.redirectUri);
    url.searchParams.set("response_type", "code");
    url.searchParams.set("scope", "openid profile email");
    url.searchParams.set("state", state);
    url.searchParams.set("nonce", nonce);
    url.searchParams.set("code_challenge", challenge);
    url.searchParams.set("code_challenge_method", "S256");

    return res.redirect(url.toString());
  } catch (error) {
    return res.status(503).send(error instanceof Error ? error.message : String(error));
  }
}

export async function finishWebLogin(req: Request, res: Response) {
  try {
    const cfg = webConfig(req);
    const code = String(req.query.code ?? "");
    const state = String(req.query.state ?? "");
    const expectedState = readCookie(req, STATE_COOKIE);
    const verifier = readCookie(req, VERIFIER_COOKIE);
    const nonce = readCookie(req, NONCE_COOKIE);
    const returnTo = safeReturnTo(readCookie(req, RETURN_COOKIE));

    if (!code || !state || !expectedState || state !== expectedState || !verifier || !nonce) {
      return res.status(400).send("Invalid or expired login state.");
    }

    const tokenResponse = await fetch(
      `${cfg.issuer}/protocol/openid-connect/token`,
      {
        method: "POST",
        headers: { "content-type": "application/x-www-form-urlencoded" },
        body: new URLSearchParams({
          grant_type: "authorization_code",
          client_id: cfg.clientId,
          client_secret: cfg.clientSecret,
          code,
          redirect_uri: cfg.redirectUri,
          code_verifier: verifier
        }),
        signal: AbortSignal.timeout(15000)
      }
    );

    if (!tokenResponse.ok) {
      const detail = (await tokenResponse.text()).slice(0, 500);
      return res.status(401).send(`Login token exchange failed: ${detail}`);
    }

    const tokens = await tokenResponse.json() as {
      id_token?: string;
      access_token?: string;
    };

    if (!tokens.id_token) {
      return res.status(401).send("Keycloak did not return an ID token.");
    }

    if (!webJwks) {
      webJwks = createRemoteJWKSet(
        new URL(`${cfg.issuer}/protocol/openid-connect/certs`)
      );
    }

    const { payload } = await jwtVerify(tokens.id_token, webJwks, {
      issuer: cfg.issuer,
      audience: cfg.clientId
    });

    if (payload.nonce !== nonce) {
      return res.status(401).send("OIDC nonce validation failed.");
    }

    const session = await createSession({
      sub: payload.sub,
      email: payload.email,
      preferred_username: payload.preferred_username,
      name: payload.name
    });

    res.append("Set-Cookie", cookie(SESSION_COOKIE, session, 8 * 60 * 60));
    for (const name of [STATE_COOKIE, VERIFIER_COOKIE, NONCE_COOKIE, RETURN_COOKIE]) {
      res.append("Set-Cookie", clearCookie(name));
    }

    return res.redirect(returnTo);
  } catch (error) {
    return res.status(401).send(error instanceof Error ? error.message : String(error));
  }
}

export async function requireWebAuth(req: Request, res: Response, next: NextFunction) {
  const token = readCookie(req, SESSION_COOKIE);
  if (!token) {
    return res.redirect(`/auth/login?returnTo=${encodeURIComponent(req.originalUrl || "/")}`);
  }

  try {
    (req as any).user = await verifySession(token);
    return next();
  } catch {
    res.append("Set-Cookie", clearCookie(SESSION_COOKIE));
    return res.redirect(`/auth/login?returnTo=${encodeURIComponent(req.originalUrl || "/")}`);
  }
}

export async function logoutWeb(_req: Request, res: Response) {
  res.append("Set-Cookie", clearCookie(SESSION_COOKIE));
  return res.redirect("/");
}