@router.get("/{provider}/login")
async def login(provider: str,redirect_uri: str | None = Query(None),set_cookie : bool = Query(True)):
    """
        Performs the login to the IDP using their authorisation endpoint. It then redirects to the token exchange endpoint.
        Example: https://localhost:8000/auth/linkedin/login
    """

    idp = get_provider(provider)

    # Generate secure state
   
    if(redirect_uri):
        #Check on allowed redirects list
        if redirect_uri not in ALLOWED_REDIRECTS :
            print("REDIRECT NOT AUTHORIZED")
            return RedirectResponse(f"{redirect_uri}?error=unauthorised")

    #The redirect uri is encoded in the state data, this is for redirects from different clients, e.g. android, web
    state_data = {
        "csrf": str(uuid.uuid4()),
        "redirect_uri": redirect_uri,
        "set_cookie" : set_cookie
    }
    state = base64.urlsafe_b64encode(json.dumps(state_data).encode()).decode()

    auth_url = await idp.get_auth_url(state)

    # Redirect browser immediately
    response = RedirectResponse(auth_url)

    # Store state in a cookie
    response.set_cookie(
        key=f"oauth_state_{provider}",
        value=state,
        httponly=True,
        secure=True,
        samesite="none"
    )

    return response
                            
                        

                            
@router.get("/{provider}/callback", response_model=str)
async def auth_callback_with_redirect(request: Request, provider: str, code: str, state: str | None = Query(None)):
    """
        Handles the callback from the IDP
        1. Takes the code from the payload and exchanges it for a token
        2. The token is verified and the user profile data returned
        3. User profile data is stored in the database
        4. JWT token is issued and set within the session cookie
    """

    idp = get_provider(provider)

    #Handle the state if it is in the payload
    stored_state = request.cookies.get(f"oauth_state_{provider}")
    if state:
        print("STATES", stored_state, state)
        if stored_state != state:
            raise HTTPException(status_code=401,detail="Invalid state")
    #Check for redirect URI in stored_state
    state_data = json.loads(base64.urlsafe_b64decode(stored_state).decode())
    redirect_uri = state_data.get("redirect_uri")
    set_cookie = state_data.get("set_cookie")
    print("REDIRECT URI IS ", redirect_uri, set_cookie)

    #State is passed in as some providers need to pass it to the token endpoint
    access_token = await idp.exchange_code(code, state)

    #Verify the token and return the user profile data
    user_profile = await idp.get_user_info(access_token)

    # database logic here
    user_record = await get_or_add_user(str(user_profile["id"]),provider,None,user_profile['email'])
    if not user_record:
        raise HTTPException(
            status_code=400,
            detail="Failed to create the user"
        )
    #Issue a JWT
    jwt_token_pair = obtain_jwt_pair(str(user_record["id"]),user_record["idp"], user_record["alias"], user_record["terms_accepted"]) 
 
    #Set the redirect URI depending on whether it exists in the cookie set to default if not in cookie
    response_redirect_uri = redirect_uri if redirect_uri else os.environ.get("CLIENT_REDIRECT_URI")

    #Divert flow depending on whether delivering tokens in cookie or sending auth code
    if set_cookie:
        response = RedirectResponse(
            url=response_redirect_uri,
            status_code=302
        )
        # Access token cookie
        response.set_cookie(
            key="access_token",
            value=jwt_token_pair["access"],
            httponly=True,
            secure=True,          # HTTPS only
            samesite="none",
            max_age=ACCESS_TOKEN_LIFETIME,
        )

        # Refresh token cookie
        response.set_cookie(
            key="refresh_token",
            value=jwt_token_pair["refresh"],
            httponly=True,
            secure=True,
            samesite="none",
            max_age=REFRESH_TOKEN_LIFETIME, 
        )
    else :
        #Generate an auth code
        auth_code = await create_auth_code(user_record["id"])
        response = RedirectResponse(
            url= f"{response_redirect_uri}?code={auth_code.code}",
            status_code=302
        )

    return response

                            
                        

                            
@router.post("/exchangeauthcodeforjwt", response_model=TokenSchema)
async def exchange_auth_code_for_jwt(auth_code : AuthCodeSchema):
    """
        This endpoint is called by the client to exchange an auth code for a JWT token and return the response as Json
        This is designed to support authorisation flows where there isn't an option of setting an authorisaton cookie
    """
    #Get the user from the database if the code is valid
    user = await validate_auth_code(auth_code.auth_code)
    if not user:
        raise HTTPException(status_code=401,detail="Authorisation code is invalid")
    
    jwt_token_pair = obtain_jwt_pair(user.id, user.idp, user.alias, user.terms_accepted)

    response = TokenSchema(
        access_token = jwt_token_pair['access'],
        refresh_token = jwt_token_pair['refresh'],
    )

    return response

                            
                        

                            
@router.get("/session", response_model=UserProfileSchema)
async def get_session(token_data = Depends(validate_jwt)):
    response = UserProfileSchema(
        id=token_data["user_id"],
        idp= token_data["idp"],
        accepted_terms = token_data["accepted_terms"],
        alias=token_data["alias"]
    )
    return response
                            
                        

                            
@router.post("/refresh")
async def refresh_jwt(request: Request, response: Response, refresh: Optional[RefreshTokenSchema], set_cookie : bool = Query(True)):
    """
        Takes the refresh token and issues a new token pair and sets the session cookie
    """

    # Fallback to cookie (browser clients)
    if not refresh.token:
        refresh_token = request.cookies.get("refresh_token")
    else:
        refresh_token = refresh.token
    try:
        jwt_token_pair = refresh_jwt_pair(refresh_token)

    except RefreshTokenExpiredError as e:
        print("Refresh token expired", e)
        raise HTTPException(status_code=401, detail="Refresh token expired")

    except InvalidRefreshTokenError as e:
        print("Refresh token invalid", e)
        raise HTTPException(status_code=401, detail="Refresh token invalid")

    # Set new cookies
    if set_cookie:
        response.set_cookie(
            key="access_token",
            value=jwt_token_pair["access"],
            httponly=True,
            secure=True,
            samesite="none",
            max_age=ACCESS_TOKEN_LIFETIME
        )

        response.set_cookie(
            key="refresh_token",
            value=jwt_token_pair["refresh"],
            httponly=True,
            secure=True,
            samesite="none",
            max_age=REFRESH_TOKEN_LIFETIME
        )

        return {"status": "refreshed"}
    
    token_pair = TokenSchema(
        access_token = jwt_token_pair['access'],
        refresh_token = jwt_token_pair['refresh'], 
    )

    return token_pair
                            
                        

                            
async def validate_jwt(request: Request):

    token = None

    # 1. Try Authorization header (mobile/API clients)
    auth_header = request.headers.get("Authorization")
    if auth_header and auth_header.startswith("Bearer "):
        token = auth_header.split(" ")[1]

    # 2. Fallback to cookie (browser clients)
    if not token:
        token = request.cookies.get("access_token")

    if not token:
        raise HTTPException(status_code=401, detail="Authentication token missing.")

    try:
        payload = jwt.decode(token, os.environ.get("SECRET_KEY"), algorithms=[ALGORITHM])
    except ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Token has expired.")
    except InvalidAudienceError:
        raise HTTPException(status_code=401, detail="Invalid audience.")
    except InvalidSignatureError:
        raise HTTPException(status_code=401, detail="Invalid token signature.")
    except Exception:
        raise HTTPException(status_code=401, detail="Invalid token.")

    return payload
                            
                        

                            
def refresh_jwt_pair(refresh_token: str):
    try:
        payload = jwt.decode(
            refresh_token,
            os.environ.get("SECRET_KEY"),
            algorithms=[ALGORITHM]
        )

        # Ensure this is a refresh token
        if payload.get("type") != "refresh":
            raise Exception("Invalid token type")

        user_id = payload.get("user_id")

        # If you stored these in the refresh token (recommended)
        alias = payload.get("alias")
        idp = payload.get("idp")
        accepted_terms = payload.get("accepted_terms")

        # Generate new tokens
        return obtain_jwt_pair(user_id, idp, alias, accepted_terms)

    except jwt.ExpiredSignatureError:
        raise RefreshTokenExpiredError("Refresh token expired")

    except jwt.InvalidTokenError:
        raise InvalidRefreshTokenError("Invalid refresh token")
    
                            
                        

                            
@router.post("/logout")
def logout(response: Response):
    """
        Logs out the user by clearing the Cookies
    """
    response.delete_cookie(
        key="access_token",
        path="/",
        secure=True,
        samesite="none",
    )

    response.delete_cookie(
        key="refresh_token",
        path="/",
        secure=True,
        samesite="none",
    )
    return {"status": "logged_out"}
                            
                        

                            
@router.post("/acceptterms", response_model=UserProfileSchema)
async def accept_terms(response: Response, set_cookie : bool = Query(True), token_data = Depends(validate_jwt)):
    """
        Accepts the terms and conditions in the database and then updates the token
    """
    try:
        await update_terms_accepted(token_data["user_id"])
    except Exception as e:
        raise HTTPException(status_code=400, detail="Unable to update the terms and conditions.")
    
    #Issue a new JWT with the updated accepted terms
    jwt_token_pair = obtain_jwt_pair(token_data["user_id"], token_data["idp"], token_data["alias"], True)
    
    # Set new cookies
    if set_cookie:
        response.set_cookie(
            key="access_token",
            value=jwt_token_pair["access"],
            httponly=True,
            secure=True,
            samesite="none",
            max_age=ACCESS_TOKEN_LIFETIME
        )

        response.set_cookie(
            key="refresh_token",
            value=jwt_token_pair["refresh"],
            httponly=True,
            secure=True,
            samesite="none",
            max_age=REFRESH_TOKEN_LIFETIME
        )

    return UserProfileSchema(
        id=token_data["user_id"],
        idp= token_data["idp"],
        accepted_terms = True,
        alias=token_data["alias"]
    )
                            
                        

                            
  useEffect(() => {
    getSession().then((res) => {
      //Set loaded, profile and authenticated 
      setLoadingSession(false);
      setAuthenticated(true);
      setProfile(res);
    }).catch(err => {
      if(err.response.status === 401) {
        //Try refresh
        refresh().then( async (res) => {
            const profile = await getSession();
            setProfile(profile);
            setLoadingSession(false);
            setAuthenticated(true);
        }).catch(async err => {
          if(err.response.status === 401) {
            //Re-authentication required
            setLoadingSession(false);
            setAuthenticated(false);
            //Authentication has failed so the user needs to sign in
            try {
              const idpList = await getIdpList();
              setIdps(idpList);
            } catch(err) {
              console.error("An error occurred retrieving the IDP list.", err);
            }

          }
        });
      }
    });
  }, [sessionRefresh])