import { Context, Effect, Layer } from "effect";
import * as jose from "jose";
import { createHash } from "node:crypto";
import { AuthError, JWTPayload } from "../../../shared/types.js";

const TOKEN_EXPIRY_DAYS = 7;
const ALGORITHM = "RS256";

let keyPair: { publicKey: jose.CryptoKey; privateKey: jose.CryptoKey } | null = null;

export class TokenService extends Context.Tag("TokenService")<
  TokenService,
  {
    readonly generateToken: (
      payload: Omit<JWTPayload, "iat" | "exp">
    ) => Effect.Effect<string, AuthError>;
    readonly verifyToken: (
      token: string
    ) => Effect.Effect<JWTPayload, AuthError>;
    readonly hashToken: (token: string) => Effect.Effect<string, never>;
  }
>() {
  static Live = Layer.effect(
    this,
    Effect.gen(function* (_) {
      // Generate key pair if not already generated
      if (!keyPair) {
        const { publicKey, privateKey } = yield* _(
          Effect.tryPromise({
            try: () => jose.generateKeyPair(ALGORITHM),
            catch: (error) =>
              new AuthError(
                "KEY_GENERATION_FAILED",
                "Failed to generate RSA key pair",
                error
              ),
          })
        );
        keyPair = { publicKey, privateKey };
      }

      const generateToken = (
        payload: Omit<JWTPayload, "iat" | "exp">
      ): Effect.Effect<string, AuthError> =>
        Effect.gen(function* (_) {
          const currentKeyPair = keyPair;
          if (!currentKeyPair) {
            return yield* _(
              Effect.fail(
                new AuthError(
                  "KEY_NOT_INITIALIZED",
                  "Key pair not initialized"
                )
              )
            );
          }

          const now = Math.floor(Date.now() / 1000);

          // Type assertion needed for index signature compatibility
          const sub = payload["sub"];
          const username = payload["username"];
          const email = payload["email"];
          const jti = payload["jti"];

          if (
            typeof sub !== "string" ||
            typeof username !== "string" ||
            typeof email !== "string" ||
            typeof jti !== "string"
          ) {
            return yield* _(
              Effect.fail(
                new AuthError(
                  "INVALID_PAYLOAD",
                  "Payload contains invalid field types"
                )
              )
            );
          }

          const fullPayload: JWTPayload = {
            sub,
            username,
            email,
            jti,
            iat: now,
            exp: now + TOKEN_EXPIRY_DAYS * 24 * 60 * 60,
          };

          return yield* _(
            Effect.tryPromise({
              try: () =>
                new jose.SignJWT(fullPayload)
                  .setProtectedHeader({ alg: ALGORITHM })
                  .setIssuedAt(fullPayload.iat)
                  .setExpirationTime(fullPayload.exp)
                  .sign(currentKeyPair.privateKey),
              catch: (error) =>
                new AuthError(
                  "TOKEN_GENERATION_FAILED",
                  "Failed to generate JWT token",
                  error
                ),
            })
          );
        });

      const verifyToken = (token: string): Effect.Effect<JWTPayload, AuthError> =>
        Effect.gen(function* (_) {
          const currentKeyPair = keyPair;
          if (!currentKeyPair) {
            return yield* _(
              Effect.fail(
                new AuthError(
                  "KEY_NOT_INITIALIZED",
                  "Key pair not initialized"
                )
              )
            );
          }

          const result = yield* _(
            Effect.tryPromise({
              try: () => jose.jwtVerify(token, currentKeyPair.publicKey),
              catch: (error) =>
                new AuthError(
                  "TOKEN_VERIFICATION_FAILED",
                  "Invalid or expired token",
                  error
                ),
            })
          );

          const payload = result.payload;

          // Validate payload structure using bracket notation for index signature properties
          if (
            typeof payload.sub !== "string" ||
            typeof payload["username"] !== "string" ||
            typeof payload["email"] !== "string" ||
            typeof payload.iat !== "number" ||
            typeof payload.exp !== "number" ||
            typeof payload.jti !== "string"
          ) {
            return yield* _(
              Effect.fail(
                new AuthError(
                  "INVALID_TOKEN_PAYLOAD",
                  "Token payload is missing required fields"
                )
              )
            );
          }

          return {
            sub: payload.sub,
            username: payload["username"],
            email: payload["email"],
            iat: payload.iat,
            exp: payload.exp,
            jti: payload.jti,
          };
        });

      const hashToken = (token: string): Effect.Effect<string, never> =>
        Effect.sync(() => {
          return createHash("sha256").update(token).digest("hex");
        });

      return { generateToken, verifyToken, hashToken };
    })
  );
}