"use client";

import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useState,
  ReactNode,
} from "react";

const API_BASE_URL =
  process.env.NEXT_PUBLIC_API_BASE_URL ?? "http://localhost:8080";

type AuthUser = {
  uuid: string;
  email: string;
  displayName?: string | null;
  role: string;
} | null;

type AuthContextValue = {
  user: AuthUser;
  token: string | null;
  loading: boolean;
  login: (params: { email: string; password: string }) => Promise<void>;
  register: (params: {
    email: string;
    password: string;
    displayName?: string;
  }) => Promise<void>;
  logout: () => void;
  getAuthHeaders: () => HeadersInit;

  /**
   * Used for non-password auth flows (ex: magic link).
   * Persists session exactly like login/register.
   */
  setSession: (token: string, user: NonNullable<AuthUser>) => void;
};

const AuthContext = createContext<AuthContextValue | undefined>(undefined);

const TOKEN_KEY = "ballistic_auth_token";
const USER_KEY = "ballistic_auth_user";

export function AuthProvider({ children }: { children: ReactNode }) {
  const [token, setToken] = useState<string | null>(null);
  const [user, setUser] = useState<AuthUser>(null);
  const [loading, setLoading] = useState<boolean>(true);

  const persistAuth = useCallback(
    (nextToken: string | null, nextUser: AuthUser) => {
      if (typeof window === "undefined") return;

      if (nextToken) {
        window.localStorage.setItem(TOKEN_KEY, nextToken);
      } else {
        window.localStorage.removeItem(TOKEN_KEY);
      }

      if (nextUser) {
        window.localStorage.setItem(USER_KEY, JSON.stringify(nextUser));
      } else {
        window.localStorage.removeItem(USER_KEY);
      }
    },
    []
  );

  // Hydrate from localStorage on first load
  useEffect(() => {
    if (typeof window === "undefined") return;

    const storedToken = window.localStorage.getItem(TOKEN_KEY);
    const storedUser = window.localStorage.getItem(USER_KEY);

    if (storedToken) setToken(storedToken);

    if (storedUser) {
      try {
        setUser(JSON.parse(storedUser));
      } catch {
        window.localStorage.removeItem(USER_KEY);
      }
    }

    setLoading(false);
  }, []);

  const setSession = useCallback(
    (nextToken: string, nextUser: NonNullable<AuthUser>) => {
      setToken(nextToken);
      setUser(nextUser);
      persistAuth(nextToken, nextUser);
    },
    [persistAuth]
  );

  const login = useCallback(
    async ({ email, password }: { email: string; password: string }) => {
      setLoading(true);
      try {
        const res = await fetch(`${API_BASE_URL}/api/auth/login`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ email, password }),
        });

        if (!res.ok) {
          const text = await res.text().catch(() => res.statusText);
          throw new Error(text || "Login failed");
        }

        const data = await res.json();

        const nextToken: string = data.token ?? data.accessToken;
        const nextUser: AuthUser =
          data.user ??
          ({
            email: data.email ?? email,
            displayName: data.displayName ?? null,
            role: data.role ?? "USER",
          } as AuthUser);

        setSession(nextToken, nextUser as NonNullable<AuthUser>);
      } finally {
        setLoading(false);
      }
    },
    [setSession]
  );

  const register = useCallback(
    async ({
      email,
      password,
      displayName,
    }: {
      email: string;
      password: string;
      displayName?: string;
    }) => {
      setLoading(true);
      try {
        const res = await fetch(`${API_BASE_URL}/api/auth/register`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ email, password, displayName }),
        });

        if (!res.ok) {
          const text = await res.text().catch(() => res.statusText);
          throw new Error(text || "Registration failed");
        }

        const data = await res.json();

        const nextToken: string = data.token ?? data.accessToken;
        const nextUser: AuthUser =
          data.user ??
          ({
            email: data.email ?? email,
            displayName: data.displayName ?? displayName ?? null,
            role: data.role ?? "USER",
          } as AuthUser);

        setSession(nextToken, nextUser as NonNullable<AuthUser>);
      } finally {
        setLoading(false);
      }
    },
    [setSession]
  );

  const logout = useCallback(() => {
    setToken(null);
    setUser(null);
    persistAuth(null, null);
  }, [persistAuth]);

  const getAuthHeaders = useCallback((): HeadersInit => {
    return token ? { Authorization: `Bearer ${token}` } : {};
  }, [token]);

  const value: AuthContextValue = {
    user,
    token,
    loading,
    login,
    register,
    logout,
    getAuthHeaders,
    setSession, // ✅ important
  };

  return <AuthContext.Provider value={value}>{children}</AuthContext.Provider>;
}

export function useAuth(): AuthContextValue {
  const ctx = useContext(AuthContext);
  if (!ctx) {
    throw new Error("useAuth must be used within an AuthProvider");
  }
  return ctx;
}